go_mod_proxy_failover.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"strings"

	"github.com/goproxy/goproxy"
)

const (
	oldDomain = "gitlab.olddomain.net"
	newDomain = "gitlab.newdomain.net"
)

// domainRewriteRoundTripper is a custom http.RoundTripper that rewrites
// URLs and response bodies for gitlab domain changes
type domainRewriteRoundTripper struct {
	base http.RoundTripper
}

func (rt *domainRewriteRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	// Only rewrite URLs for the old domain
	if strings.Contains(req.URL.String(), oldDomain) {
		// Create a copy of the request to modify
		reqCopy := req.Clone(req.Context())
		
		// Rewrite the URL from old domain to new domain
		oldURL := reqCopy.URL.String()
		newURL := strings.Replace(oldURL, oldDomain, newDomain, -1)
		
		var err error
		reqCopy.URL, err = req.URL.Parse(newURL)
		if err != nil {
			return nil, fmt.Errorf("failed to parse rewritten URL: %v", err)
		}
		
		// Rewrite any Host header if present
		if reqCopy.Host != "" && strings.Contains(reqCopy.Host, oldDomain) {
			reqCopy.Host = strings.Replace(reqCopy.Host, oldDomain, newDomain, -1)
		}
		
		// If the request has a body, rewrite its content too
		if reqCopy.Body != nil {
			body, err := io.ReadAll(reqCopy.Body)
			if err != nil {
				return nil, fmt.Errorf("failed to read request body: %v", err)
			}
			reqCopy.Body.Close()
			
			// Replace old domain with new domain in the body
			newBody := strings.Replace(string(body), oldDomain, newDomain, -1)
			reqCopy.Body = io.NopCloser(bytes.NewReader([]byte(newBody)))
			reqCopy.ContentLength = int64(len(newBody))
		}
		
		// Use the modified request
		req = reqCopy
	}
	
	// Perform the HTTP request with the base transport
	resp, err := rt.base.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	
	// For responses, we want to rewrite from new domain back to old domain
	// to maintain compatibility with existing import paths
	if strings.Contains(req.URL.String(), newDomain) {
		if resp.Body != nil {
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				resp.Body.Close()
				return nil, fmt.Errorf("failed to read response body: %v", err)
			}
			resp.Body.Close()
			
			// Replace new domain with old domain in the response body
			newBody := strings.Replace(string(body), newDomain, oldDomain, -1)
			resp.Body = io.NopCloser(bytes.NewReader([]byte(newBody)))
			resp.ContentLength = int64(len(newBody))
			
			// Update content length header
			resp.Header.Set("Content-Length", fmt.Sprintf("%d", len(newBody)))
		}
	}
	
	return resp, nil
}

func main() {
	// Get port from environment or use default
	port := os.Getenv("PORT")
	if port == "" {
		port = "8081"
	}
	
	// Configure the goproxy instance
	g := &goproxy.Goproxy{
		Cacher: goproxy.DirCacher("cache"),
	}
	
	// Create our custom transport that handles domain rewriting
	transport := &domainRewriteRoundTripper{
		base: http.DefaultTransport,
	}
	
	// Set our custom transport as the goproxy's transport
	g.Transport = transport
	
	// Setup logger
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	g.Logger = logger
	
	// Print startup message
	logger.Info("Starting Go Module Proxy with domain rewriting", 
		"port", port,
		"oldDomain", oldDomain,
		"newDomain", newDomain,
	)
	
	// Start the server
	if err := http.ListenAndServe(":"+port, g); err != nil {
		logger.Error("Failed to start server", "error", err)
		os.Exit(1)
	}
}