proxy_server.go
1//go:build ignore
2// +build ignore
3
4// This example demonstrates how to use the ysweet proxy handler to create
5// a websocket proxy server that forwards connections to a y-sweet server.
6//
7// Usage:
8//
9// go run proxy_server.go
10//
11// Then connect from a Yjs client to ws://localhost:3000/d/my-doc/ws/my-doc
12package main
13
14import (
15 "context"
16 "log"
17 "net/http"
18 "os"
19 "os/signal"
20 "syscall"
21 "time"
22
23 "git.kilimanjaro.io/ygo/ysweet"
24 "github.com/coder/websocket"
25)
26
27func main() {
28 // Configuration
29 upstreamURL := getEnv("Y_SWEET_URL", "ws://localhost:8080")
30 listenAddr := getEnv("LISTEN_ADDR", ":3000")
31 authToken := getEnv("Y_SWEET_AUTH_TOKEN", "")
32
33 // Create a master context for graceful shutdown
34 // When this context is cancelled, all active websocket connections
35 // will be closed with a 5-second drain period
36 masterCtx, cancel := context.WithCancel(context.Background())
37 defer cancel()
38
39 // Example 1: Use default URL pattern (/d/{docID}/ws/{docID})
40 defaultHandler := ysweet.ProxyHandler(upstreamURL,
41 ysweet.WithContext(masterCtx), // Enable graceful shutdown
42 ysweet.WithTargetAuthToken(authToken),
43 ysweet.WithOnConnect(func(ctx context.Context, docID string, w http.ResponseWriter, r *http.Request) error {
44 log.Printf("[CONNECT] Client %s connecting to doc %s", r.RemoteAddr, docID)
45 // You can now take action based on docID, e.g., check permissions
46 return nil
47 }),
48 ysweet.WithOnWebsocketUpgrade(func(ctx context.Context, docID string, clientConn, upstreamConn *websocket.Conn) error {
49 log.Printf("[UPGRADE] Doc %s websocket upgrade successful", docID)
50 return nil
51 }),
52 ysweet.WithOnDisconnect(func(ctx context.Context, docID string, reason ysweet.DisconnectReason) error {
53 log.Printf("[DISCONNECT] Doc %s disconnected, reason: %v", docID, reason)
54 return nil
55 }),
56 )
57
58 // Example 2: Custom URL pattern using Go 1.22+ path parameters
59 customHandler := ysweet.ProxyHandler(upstreamURL,
60 ysweet.WithContext(masterCtx), // Enable graceful shutdown
61 ysweet.WithDocIDFunc(func(r *http.Request) string {
62 // Extract docID from custom URL: /ws/{docID}
63 return r.PathValue("docID")
64 }),
65 ysweet.WithOnConnect(func(ctx context.Context, docID string, w http.ResponseWriter, r *http.Request) error {
66 log.Printf("[CONNECT] Custom handler - doc %s connecting", docID)
67 return nil
68 }),
69 )
70
71 // Set up HTTP routes
72 mux := http.NewServeMux()
73
74 // Default y-sweet pattern
75 mux.Handle("/d/", defaultHandler)
76
77 // Custom pattern: /ws/{docID} -> proxies to upstream /d/{docID}/ws/{docID}
78 mux.Handle("/ws/{docID}", customHandler)
79
80 // Health check endpoint
81 mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
82 w.WriteHeader(http.StatusOK)
83 w.Write([]byte("ok"))
84 })
85
86 // Create server
87 server := &http.Server{
88 Addr: listenAddr,
89 Handler: mux,
90 }
91
92 // Handle graceful shutdown
93 sigChan := make(chan os.Signal, 1)
94 signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
95
96 go func() {
97 <-sigChan
98 log.Println("Shutting down - closing all websocket connections...")
99
100 // Cancel master context first to close all websocket connections
101 // The proxy will wait 5 seconds for connections to drain to upstream
102 cancel()
103
104 // Then shutdown the HTTP server
105 shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
106 defer shutdownCancel()
107
108 if err := server.Shutdown(shutdownCtx); err != nil {
109 log.Printf("Server shutdown error: %v", err)
110 }
111 }()
112
113 log.Printf("Starting proxy server on %s", listenAddr)
114 log.Printf("Proxying to y-sweet at %s", upstreamURL)
115
116 if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
117 log.Fatalf("Server error: %v", err)
118 }
119}
120
121func getEnv(key, defaultValue string) string {
122 if value := os.Getenv(key); value != "" {
123 return value
124 }
125 return defaultValue
126}
