Use proper TokenSource for IAP auth

iap/dialopts.go

@@ -1,19 +1,23 @@
 package iap
 
+import (
+	"golang.org/x/oauth2"
+)
+
 type DialOption func(*dialOptions)
 
 type dialOptions struct {
-	Zone      string
-	Token     string
-	Region    string
-	Project   string
-	Port      string
-	Network   string
-	Interface string
-	Instance  string
-	Host      string
-	Group     string
-	Compress  bool
+	Zone        string
+	TokenSource *oauth2.TokenSource
+	Region      string
+	Project     string
+	Port        string
+	Network     string
+	Interface   string
+	Instance    string
+	Host        string
+	Group       string
+	Compress    bool
 }
 
 func (d *dialOptions) collectOpts(opts []DialOption) {
@@ -22,10 +26,10 @@ func (d *dialOptions) collectOpts(opts []DialOption) {
 	}
 }
 
-// WithToken is a functional option that sets the authorization token.
-func WithToken(token string) func(*dialOptions) {
+// WithTokenSource is a functional option that sets the authorization toke source.
+func WithTokenSource(tokenSource *oauth2.TokenSource) func(*dialOptions) {
 	return func(d *dialOptions) {
-		d.Token = token
+		d.TokenSource = tokenSource
 	}
 }
 

iap/iap.go

@@ -92,11 +92,20 @@ func Dial(ctx context.Context, opts ...DialOption) (*Conn, error) {
 	dopts := &dialOptions{}
 	dopts.collectOpts(opts)
 
+	header := make(http.Header)
+	header.Set("Origin", proxyOrigin)
+
+	if dopts.TokenSource != nil {
+		token, err := (*dopts.TokenSource).Token()
+		if err != nil {
+			return nil, err
+		}
+
+		header.Set("Authorization", fmt.Sprintf("%v %v", token.Type(), token.AccessToken))
+	}
+
 	wsOptions := websocket.DialOptions{
-		HTTPHeader: http.Header{
-			"Authorization": []string{fmt.Sprintf("Bearer %v", dopts.Token)},
-			"Origin":        []string{proxyOrigin},
-		},
+		HTTPHeader:      header,
 		Subprotocols:    []string{proxySubproto},
 		CompressionMode: websocket.CompressionDisabled,
 	}

internal/cmd/root.go

@@ -1,8 +1,12 @@
 package cmd
 
 import (
+	"context"
+
 	"github.com/charmbracelet/log"
 	"github.com/spf13/cobra"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
 )
 
 var (
@@ -17,6 +21,14 @@ var rootCmd = &cobra.Command{
 	Long: "Utility for Google Cloud's Identity-Aware Proxy",
 }
 
+func getTokenSource() *oauth2.TokenSource {
+	tokenSource, err := google.DefaultTokenSource(context.Background())
+	if err != nil {
+		log.Fatal(err)
+	}
+	return &tokenSource
+}
+
 func init() {
 	rootCmd.PersistentFlags().BoolVarP(&compress, "compress", "c", false, "Enable WebSocket compression")
 	rootCmd.PersistentFlags().StringVarP(&listen, "listen", "l", "127.0.0.1:0", "Listen address and port")

internal/cmd/to_host.go

@@ -20,13 +20,14 @@ var hostCmd = &cobra.Command{
 	Long: "Create a tunnel to a remote private IP or FQDN (requires BeyondCorp Enterprise)",
 	Args: cobra.ExactArgs(1),
 	PreRun: func(cmd *cobra.Command, args []string) {
-		log.Info("Starting proxy", "listen", listen, "dest", fmt.Sprintf("%v:%v", args[0], port), "project", project)
+		log.Info("Started proxy", "listen", listen, "dest", fmt.Sprintf("%v:%v", args[0], port), "project", project)
 	},
 	Run: func(cmd *cobra.Command, args []string) {
 		opts := []iap.DialOption{
 			iap.WithProject(project),
 			iap.WithHost(args[0], region, network, destGroup),
 			iap.WithPort(fmt.Sprint(port)),
+			iap.WithTokenSource(getTokenSource()),
 		}
 		if compress {
 			opts = append(opts, iap.WithCompression())

internal/cmd/to_instance.go

@@ -26,6 +26,7 @@ var instanceCmd = &cobra.Command{
 			iap.WithProject(project),
 			iap.WithInstance(args[0], zone, ninterface),
 			iap.WithPort(fmt.Sprint(port)),
+			iap.WithTokenSource(getTokenSource()),
 		}
 		if compress {
 			opts = append(opts, iap.WithCompression())

internal/proxy/proxy.go

@@ -7,14 +7,12 @@ import (
 
 	"github.com/cedws/iapc/iap"
 	"github.com/charmbracelet/log"
-	"golang.org/x/oauth2/google"
 )
 
 // Start starts a proxy server that listens on the given address and port.
 func Start(listen string, opts []iap.DialOption) {
-	opts = append(opts, iap.WithToken(getToken()))
 	if err := testConn(opts); err != nil {
-		log.Fatal(err)
+		log.Fatalf("Error testing connection: %v", err)
 	}
 
 	listener, err := net.Listen("tcp", listen)
@@ -45,7 +43,7 @@ func handleClient(opts []iap.DialOption, conn net.Conn) {
 
 	tun, err := iap.Dial(context.Background(), opts...)
 	if err != nil {
-		log.Error(err)
+		log.Errorf("Error dialing IAP: %v", err)
 		return
 	}
 	defer tun.Close()
@@ -54,17 +52,5 @@ func handleClient(opts []iap.DialOption, conn net.Conn) {
 	go io.Copy(conn, tun)
 	io.Copy(tun, conn)
 
-	log.Info("Client disconnected", "client", conn.RemoteAddr())
-}
-
-func getToken() string {
-	credentials, err := google.FindDefaultCredentials(context.Background())
-	if err != nil {
-		log.Fatal(err)
-	}
-	tok, err := credentials.TokenSource.Token()
-	if err != nil {
-		log.Fatal(err)
-	}
-	return tok.AccessToken
+	log.Info("Client disconnected", "client", conn.RemoteAddr(), "sentbytes", tun.Sent(), "recvbytes", tun.Received())
 }