main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

var (
	listenAddr   = flag.String("listen", "0.0.0.0:11434", "Address to listen on")
	targetAddr   = flag.String("target", "127.0.0.1:11434", "Target Ollama server address")
	apiKey       = flag.String("apikey", "ollama-gateway", "API key for authentication")
	apiKeyParam  = flag.String("apikey-param", "api_key", "URL parameter name for API key")
	apiKeyHeader = flag.String("apikey-header", "X-API-Key", "Header name for API key")
	whitelistIPs = flag.String("whitelist", "", "Comma-separated list of IP addresses/ranges in CIDR notation (e.g. 127.0.0.1/32,192.168.0.1/24)")
	trustedIPs   = flag.String("trusted-ips", "", "Comma-separated list of IP addresses/ranges in CIDR notation that can bypass authentication")
	debug        = flag.Bool("debug", false, "Enable debug logging")
	certFile     = flag.String("cert", "", "Path to TLS certificate file for HTTPS mode")
	keyFile      = flag.String("key", "", "Path to TLS key file for HTTPS mode")
)

func main() {
	flag.Parse()

	fmt.Printf("Starting Ollama Gateway proxy server\n")
	fmt.Printf("Listening on: %s\n", *listenAddr)
	fmt.Printf("Forwarding to: %s\n", *targetAddr)
	fmt.Printf("API Key required: %s\n", *apiKey)
	fmt.Printf("API Key can be provided via URL parameter '%s' or header '%s'\n", *apiKeyParam, *apiKeyHeader)

	// Check if HTTPS mode is enabled
	httpsMode := *certFile != "" && *keyFile != ""
	if httpsMode {
		fmt.Printf("HTTPS mode enabled with certificate: %s and key: %s\n", *certFile, *keyFile)
	}

	// Parse IP whitelist if provided
	var ipWhitelist []*net.IPNet
	if *whitelistIPs != "" {
		fmt.Printf("IP whitelist enabled\n")
		cidrs := strings.Split(*whitelistIPs, ",")
		for _, cidr := range cidrs {
			_, ipNet, err := net.ParseCIDR(strings.TrimSpace(cidr))
			if err != nil {
				log.Fatalf("Invalid CIDR in whitelist: %s - %v", cidr, err)
			}
			ipWhitelist = append(ipWhitelist, ipNet)
			fmt.Printf("Whitelisted: %s\n", ipNet.String())
		}
	}

	// Parse trusted IPs if provided
	var trustedIPNets []*net.IPNet
	if *trustedIPs != "" {
		fmt.Printf("Trusted IPs enabled (these IPs bypass authentication)\n")
		cidrs := strings.Split(*trustedIPs, ",")
		for _, cidr := range cidrs {
			_, ipNet, err := net.ParseCIDR(strings.TrimSpace(cidr))
			if err != nil {
				log.Fatalf("Invalid CIDR in trusted IPs: %s - %v", cidr, err)
			}
			trustedIPNets = append(trustedIPNets, ipNet)
			fmt.Printf("Trusted IP: %s\n", ipNet.String())
		}
	}

	// Create a reverse proxy
	targetURL, err := url.Parse("http://" + *targetAddr)
	if err != nil {
		log.Fatalf("Error parsing target URL: %v", err)
	}

	proxy := httputil.NewSingleHostReverseProxy(targetURL)

	// Create a custom director function to modify the request
	originalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		// Call the original director
		originalDirector(req)

		// Set the host header to match the target
		req.Host = targetURL.Host

		// Remove the API key from the URL if present
		q := req.URL.Query()
		if q.Get(*apiKeyParam) != "" {
			q.Del(*apiKeyParam)
			req.URL.RawQuery = q.Encode()
		}

		// Remove the API key header if present
		req.Header.Del(*apiKeyHeader)

		// Set X-Forwarded-For header
		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			if prior, ok := req.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			req.Header.Set("X-Forwarded-For", clientIP)
		}

		if *debug {
			log.Printf("Forwarding request: %s %s", req.Method, req.URL.String())
		}
	}

	// Create a custom error handler
	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		log.Printf("Error forwarding request: %v", err)
		http.Error(w, "Error forwarding request to Ollama server", http.StatusBadGateway)
	}

	// Create a custom transport to modify the response
	proxy.ModifyResponse = func(resp *http.Response) error {
		if *debug {
			log.Printf("Response status: %d", resp.StatusCode)
		}
		return nil
	}

	// Create a handler that checks for API key before forwarding
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Extract client IP for both whitelist and trusted IP checks
		clientIP, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			log.Printf("Error parsing client IP %s: %v", r.RemoteAddr, err)
			http.Error(w, "Internal server error", http.StatusInternalServerError)
			return
		}

		ip := net.ParseIP(clientIP)
		if ip == nil {
			log.Printf("Invalid client IP: %s", clientIP)
			http.Error(w, "Internal server error", http.StatusInternalServerError)
			return
		}

		// Check if IP is in whitelist if whitelist is enabled
		if len(ipWhitelist) > 0 {
			allowed := false
			for _, ipNet := range ipWhitelist {
				if ipNet.Contains(ip) {
					allowed = true
					break
				}
			}

			if !allowed {
				log.Printf("Unauthorized request from %s - IP not in whitelist", clientIP)
				http.Error(w, "Unauthorized: IP not in whitelist", http.StatusForbidden)
				return
			}

			if *debug {
				log.Printf("Client IP %s allowed by whitelist", clientIP)
			}
		}

		// Check if IP is in trusted IPs list (can bypass authentication)
		isTrusted := false
		for _, ipNet := range trustedIPNets {
			if ipNet.Contains(ip) {
				isTrusted = true
				if *debug {
					log.Printf("Client IP %s is trusted, bypassing authentication", clientIP)
				}
				break
			}
		}

		// Only check API key if the IP is not trusted
		if !isTrusted {
			// Check for API key in URL parameters
			urlKey := r.URL.Query().Get(*apiKeyParam)

			// Check for API key in headers
			headerKey := r.Header.Get(*apiKeyHeader)

			// Validate API key
			if urlKey != *apiKey && headerKey != *apiKey {
				log.Printf("Unauthorized request from %s - no valid API key provided", r.RemoteAddr)
				http.Error(w, "Unauthorized: API key required", http.StatusUnauthorized)
				return
			}
		}

		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL.Path)

		// Forward the request to the target server
		proxy.ServeHTTP(w, r)
	})

	// Start the server with either HTTP or HTTPS based on provided certificates
	if httpsMode {
		log.Printf("Starting HTTPS server on %s", *listenAddr)
		log.Fatal(http.ListenAndServeTLS(*listenAddr, *certFile, *keyFile, nil))
	} else {
		log.Printf("Starting HTTP server on %s", *listenAddr)
		log.Fatal(http.ListenAndServe(*listenAddr, nil))
	}
}