Make proxy configurable to access ZTS in Golang

libs/go/ztsroletoken/role-token.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/url"
 	"os"
 	"sync"
 	"time"
@@ -34,6 +35,7 @@ type RoleToken interface {
 // for getting a role token. The zero-value is a valid configuration.
 type RoleTokenOptions struct {
 	BaseZTSURL string        // the base ZTS URL to use
+	ProxyURL   string        // the proxy URL for accessing ZTS
 	Role       string        // the single role for which a token is required
 	MinExpire  time.Duration // the minimum expiry of the token in (server default if zero)
 	MaxExpire  time.Duration // the maximum expiry of the token (server default if zero)
@@ -88,6 +90,16 @@ func (r *roleToken) updateRoleToken() (string, error) {
 		return "", errors.New("BaseZTSURL is empty")
 	}
 
+	var proxyURL *url.URL
+	if r.opts.ProxyURL != "" {
+		p, err := url.Parse(r.opts.ProxyURL)
+		if err != nil {
+			return "", err
+		} else {
+			proxyURL = p
+		}
+	}
+
 	r.l.Lock()
 	defer r.l.Unlock()
 
@@ -107,15 +119,25 @@ func (r *roleToken) updateRoleToken() (string, error) {
 			config.RootCAs = certPool
 		}
 
-		z = zts.NewClient(r.opts.BaseZTSURL, &http.Transport{
+		tr := http.Transport{
 			TLSClientConfig: config,
-		})
+		}
+		if proxyURL != nil {
+			tr.Proxy = http.ProxyURL(proxyURL)
+		}
+		z = zts.NewClient(r.opts.BaseZTSURL, &tr)
 	} else {
 		ntoken, err := r.tok.Value()
 		if err != nil {
 			return "", err
 		}
-		z = zts.NewClient(r.opts.BaseZTSURL, nil)
+		if proxyURL != nil {
+			z = zts.NewClient(r.opts.BaseZTSURL, &http.Transport{
+				Proxy: http.ProxyURL(proxyURL),
+			})
+		} else {
+			z = zts.NewClient(r.opts.BaseZTSURL, nil)
+		}
 		z.AddCredentials(r.opts.AuthHeader, ntoken)
 	}
 

libs/go/ztsroletoken/role-token_test.go

@@ -8,6 +8,8 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httputil"
+	"net/url"
 	"os"
 	"sync"
 	"testing"
@@ -50,18 +52,24 @@ WYjCE4hWTQzn0xtwrqrT/c337wvX48p4yk31WdXtCUA=
 // httptest.NewTLSServer uses a cert/key committed at net/http/internal
 // Reusing the same cert here, so that we can use it as the RootCA Cert in the client connection
 var LocalhostCert = []byte(`-----BEGIN CERTIFICATE-----
-MIICEzCCAXygAwIBAgIQMIMChMLGrR+QvmQvpwAU6zANBgkqhkiG9w0BAQsFADAS
+MIIDOTCCAiGgAwIBAgIQSRJrEpBGFc7tNb1fb5pKFzANBgkqhkiG9w0BAQsFADAS
 MRAwDgYDVQQKEwdBY21lIENvMCAXDTcwMDEwMTAwMDAwMFoYDzIwODQwMTI5MTYw
-MDAwWjASMRAwDgYDVQQKEwdBY21lIENvMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCB
-iQKBgQDuLnQAI3mDgey3VBzWnB2L39JUU4txjeVE6myuDqkM/uGlfjb9SjY1bIw4
-iA5sBBZzHi3z0h1YV8QPuxEbi4nW91IJm2gsvvZhIrCHS3l6afab4pZBl2+XsDul
-rKBxKKtD1rGxlG4LjncdabFn9gvLZad2bSysqz/qTAUStTvqJQIDAQABo2gwZjAO
-BgNVHQ8BAf8EBAMCAqQwEwYDVR0lBAwwCgYIKwYBBQUHAwEwDwYDVR0TAQH/BAUw
-AwEB/zAuBgNVHREEJzAlggtleGFtcGxlLmNvbYcEfwAAAYcQAAAAAAAAAAAAAAAA
-AAAAATANBgkqhkiG9w0BAQsFAAOBgQCEcetwO59EWk7WiJsG4x8SY+UIAA+flUI9
-tyC4lNhbcF2Idq9greZwbYCqTTTr2XiRNSMLCOjKyI7ukPoPjo16ocHj+P3vZGfs
-h1fIw3cSS2OolhloGw/XM6RWPWtPAlGykKLciQrBru5NAPvCMsb/I1DAceTiotQM
-fblo6RBxUQ==
+MDAwWjASMRAwDgYDVQQKEwdBY21lIENvMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
+MIIBCgKCAQEA6Gba5tHV1dAKouAaXO3/ebDUU4rvwCUg/CNaJ2PT5xLD4N1Vcb8r
+bFSW2HXKq+MPfVdwIKR/1DczEoAGf/JWQTW7EgzlXrCd3rlajEX2D73faWJekD0U
+aUgz5vtrTXZ90BQL7WvRICd7FlEZ6FPOcPlumiyNmzUqtwGhO+9ad1W5BqJaRI6P
+YfouNkwR6Na4TzSj5BrqUfP0FwDizKSJ0XXmh8g8G9mtwxOSN3Ru1QFc61Xyeluk
+POGKBV/q6RBNklTNe0gI8usUMlYyoC7ytppNMW7X2vodAelSu25jgx2anj9fDVZu
+h7AXF5+4nJS4AAt0n1lNY7nGSsdZas8PbQIDAQABo4GIMIGFMA4GA1UdDwEB/wQE
+AwICpDATBgNVHSUEDDAKBggrBgEFBQcDATAPBgNVHRMBAf8EBTADAQH/MB0GA1Ud
+DgQWBBStsdjh3/JCXXYlQryOrL4Sh7BW5TAuBgNVHREEJzAlggtleGFtcGxlLmNv
+bYcEfwAAAYcQAAAAAAAAAAAAAAAAAAAAATANBgkqhkiG9w0BAQsFAAOCAQEAxWGI
+5NhpF3nwwy/4yB4i/CwwSpLrWUa70NyhvprUBC50PxiXav1TeDzwzLx/o5HyNwsv
+cxv3HdkLW59i/0SlJSrNnWdfZ19oTcS+6PtLoVyISgtyN6DpkKpdG1cOkW3Cy2P2
++tK/tKHRP1Y/Ra0RiDpOAmqn0gCOFGz8+lqDIor/T7MTpibL3IxqWfPrvfVRHL3B
+grw/ZQTTIVjjh4JBSW3WyWgNo/ikC1lrVxzl4iPUGptxT36Cr7Zk2Bsg0XqwbOvK
+5d+NTDREkSnUbie4GeutujmX3Dsx88UiV6UY/4lHJa6I5leHUNOHahRbpbWeOfs/
+WkBKOclmOV2xlTVuPw==
 -----END CERTIFICATE-----`)
 
 type tokp struct {
@@ -90,7 +98,14 @@ func (rt *rtHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	out := struct {
 		Token      string `json:"token"`
 		ExpiryTime int64  `json:"expiryTime"`
-	}{Token: fmt.Sprintf("RT%d", rt.count), ExpiryTime: time.Now().Add(rt.expiry).Unix()}
+	}{}
+	// "X-Forwarded-For" header is automatically added if the request goes through a reverse proxy
+	if r.Header.Get("X-Forwarded-For") == "" {
+		out.Token = fmt.Sprintf("RT%d", rt.count)
+	} else {
+		out.Token = fmt.Sprintf("RT%d-%s", rt.count, r.Header.Get("X-Forwarded-For"))
+	}
+	out.ExpiryTime = time.Now().Add(rt.expiry).Unix()
 	b, _ := json.Marshal(&out)
 	w.Write(b)
 }
@@ -138,6 +153,36 @@ func TestRoleToken(t *testing.T) {
 	}
 }
 
+func TestRoleTokenWithProxy(t *testing.T) {
+	s := httptest.NewServer(&rtHandler{expiry: 1 * time.Minute})
+	defer s.Close()
+
+	sURL, err := url.Parse(s.URL)
+	if err != nil {
+		t.Fatal("failed to parse zts url", err)
+	}
+
+	p := httptest.NewServer(httputil.NewSingleHostReverseProxy(sURL))
+	defer p.Close()
+
+	tp := &tokp{}
+	e := 1 * time.Minute
+	rt := NewRoleToken(tp, "my.domain", RoleTokenOptions{
+		BaseZTSURL: s.URL,
+		ProxyURL:   p.URL,
+		MinExpire:  e,
+		MaxExpire:  e,
+	})
+
+	tok, err := rt.RoleTokenValue()
+	if err != nil {
+		t.Fatal("error getting role token", err)
+	}
+	if tok != "RT1-127.0.0.1" {
+		t.Error("invalid role token", tok)
+	}
+}
+
 func TestRoleTokenFromCert(t *testing.T) {
 	s := httptest.NewTLSServer(&rtHandler{expiry: 3 * time.Second})
 	defer s.Close()