proxy.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

type TokenManager struct {
	mu        sync.RWMutex
	token     string
	expiresAt time.Time
}

func (tm *TokenManager) GetToken() (string, error) {
	tm.mu.RLock()
	if tm.token != "" && time.Now().Before(tm.expiresAt) {
		defer tm.mu.RUnlock()
		return tm.token, nil
	}
	tm.mu.RUnlock()

	tm.mu.Lock()
	defer tm.mu.Unlock()

	if tm.token != "" && time.Now().Before(tm.expiresAt) {
		return tm.token, nil
	}

	out, err := exec.Command("gcloud", "auth", "print-identity-token").Output()
	if err != nil {
		return "", fmt.Errorf("failed to get identity token from gcloud: %w", err)
	}

	token := strings.TrimSpace(string(out))
	if token == "" {
		return "", fmt.Errorf("gcloud returned an empty identity token")
	}

	tm.token = token
	tm.expiresAt = time.Now().Add(45 * time.Minute)
	return tm.token, nil
}

func main() {
	targetFlag := flag.String("target", os.Getenv("TARGET_URL"), "Target Cloud Run URL")
	portFlag := flag.Int("port", 8080, "Local port to listen on")
	flag.Parse()

	if *targetFlag == "" {
		log.Fatal("Error: target URL is required. Provide -target=https://... or set TARGET_URL env variable.")
	}

	targetURL, err := url.Parse(*targetFlag)
	if err != nil || targetURL.Scheme == "" || targetURL.Host == "" {
		log.Fatalf("Invalid target URL '%s': %v", *targetFlag, err)
	}

	tm := &TokenManager{}

	// Verify token retrieval before starting
	log.Printf("Fetching initial identity token via gcloud...")
	if _, err := tm.GetToken(); err != nil {
		log.Printf("Warning: failed to get initial token: %v", err)
	} else {
		log.Printf("Successfully acquired identity token.")
	}

	proxy := httputil.NewSingleHostReverseProxy(targetURL)
	originalDirector := proxy.Director

	proxy.Director = func(req *http.Request) {
		originalDirector(req)
		req.Host = targetURL.Host

		// Inject Google Cloud IAM Identity Bearer Token
		token, err := tm.GetToken()
		if err != nil {
			log.Printf("Error refreshing token: %v", err)
		} else {
			req.Header.Set("Authorization", "Bearer "+token)
		}

		// Rewrite Origin to match target host so code-server CSRF protection passes
		if req.Header.Get("Origin") != "" {
			req.Header.Set("Origin", targetURL.Scheme+"://"+targetURL.Host)
		}

		req.Header.Set("X-Forwarded-Host", req.Host)
		req.Header.Set("X-Forwarded-Proto", "https")
	}

	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		log.Printf("Proxy error for %s %s: %v", r.Method, r.URL.Path, err)
		http.Error(w, fmt.Sprintf("Proxy Error: %v", err), http.StatusBadGateway)
	}

	addr := fmt.Sprintf("127.0.0.1:%d", *portFlag)
	log.Printf("==================================================================")
	log.Printf(" Cloud Run WebSocket Auth Proxy")
	log.Printf(" Target: %s", targetURL.String())
	log.Printf(" Listening on: http://%s", addr)
	log.Printf("==================================================================")

	if err := http.ListenAndServe(addr, proxy); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}