[feature] update: implementing websockets support

.vscode/launch.json

@@ -37,11 +37,11 @@
             "asRoot": true,
         },
         {
-            "name": "Run MCP Server",
+            "name": "Websocket Server",
             "type": "go",
             "request": "launch",
             "mode": "auto",
-            "program": "${workspaceFolder}/examples/mcp/",
+            "program": "${workspaceFolder}/examples/websocket/",
             "asRoot": true,
         },
     ]

examples/websocket/configs/config.yaml

@@ -0,0 +1,30 @@
+service:
+  http:
+    h3: 
+      enabled: true
+      address:
+        ip: ""
+        port: 4431
+    h1:
+      enabled: true
+      address:
+        ip: ""
+        port: 4431
+    h1_ssl:
+      enabled: true
+      address:
+        ip: ""
+        port: 4430
+  tls:
+    generate_if_missing: true
+    certificate:
+      raw: ""
+      path: cert.pem
+    key:
+      raw: ""
+      path: key.pem
+  mcp:
+    enabled: false
+    address:
+      ip: ""
+      port: 4432

examples/websocket/main.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/ayushanand18/crazyhttp/internal/constants"
+	crazyserver "github.com/ayushanand18/crazyhttp/pkg/server"
+	"github.com/ayushanand18/crazyhttp/pkg/types"
+)
+
+func main() {
+	ctx := context.Background()
+
+	server := crazyserver.NewHttpServer(ctx)
+	if err := server.Initialize(ctx); err != nil {
+		log.Fatalf("Server failed to Initialize: %v", err)
+	}
+
+	server.WebSocket("/ws-test").
+		WithOptions(types.WebSocketOption{
+			AllowedOrigins: []string{"*"},
+		}).
+		Serve(func(ctx context.Context) error {
+			reqChanel := ctx.Value(constants.WebsocketRequestChannel).(chan types.WebsocketStreamChunk)
+			respChanel := ctx.Value(constants.WebsocketResponseChannel).(chan types.WebsocketStreamChunk)
+
+			for chunk := range reqChanel {
+				fmt.Printf("Received chunk: ID=%d, Type=%d, Data=%s\n", chunk.Id, chunk.MessageType, string(chunk.Data))
+				respChanel <- types.WebsocketStreamChunk{
+					Data: []byte(fmt.Sprintf("Echo: %s", string(chunk.Data))),
+				}
+			}
+
+			return nil
+		})
+
+	if err := server.ListenAndServe(ctx); err != nil {
+		log.Fatalf("Server failed to start: %v", err)
+	}
+}

examples/websocket/main_test.go

@@ -0,0 +1,91 @@
+package main_test
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/ayushanand18/crazyhttp/internal/constants"
+	crazyserver "github.com/ayushanand18/crazyhttp/pkg/server"
+	"github.com/ayushanand18/crazyhttp/pkg/types"
+	"github.com/gorilla/websocket"
+)
+
+// waitForServer waits until TCP port is accepting connections
+func waitForServer(addr string) error {
+	for i := 0; i < 20; i++ {
+		conn, err := net.DialTimeout("tcp", addr, 100*time.Millisecond)
+		if err == nil {
+			conn.Close()
+			return nil
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	return fmt.Errorf("server not ready on %s", addr)
+}
+
+func TestUserRoute_WebsocketRequest(t *testing.T) {
+	ctx := context.Background()
+	addr := "localhost:4431"
+
+	server := crazyserver.NewHttpServer(ctx)
+	if err := server.Initialize(ctx); err != nil {
+		t.Fatalf("Initialize failed: %v", err)
+	}
+
+	// WebSocket endpoint
+	server.WebSocket("/ws-test").
+		WithOptions(types.WebSocketOption{AllowedOrigins: []string{"*"}}).
+		Serve(func(ctx context.Context) error {
+			reqCh := ctx.Value(constants.WebsocketRequestChannel).(chan types.WebsocketStreamChunk)
+			respCh := ctx.Value(constants.WebsocketResponseChannel).(chan types.WebsocketStreamChunk)
+
+			// Keep the handler alive to echo messages
+			for chunk := range reqCh {
+				respCh <- types.WebsocketStreamChunk{
+					Data: []byte(fmt.Sprintf("Echo: %s", chunk.Data)),
+				}
+			}
+			return nil
+		})
+
+	// Start server in background
+	go func() {
+		if err := server.ListenAndServe(ctx); err != nil {
+			t.Logf("Server stopped: %v", err)
+		}
+	}()
+
+	// Wait for server to be ready
+	if err := waitForServer(addr); err != nil {
+		t.Fatalf("Server not ready: %v", err)
+	}
+
+	// Connect WebSocket client
+	wsURL := fmt.Sprintf("ws://%s/ws-test", addr)
+	dialer := websocket.Dialer{}
+	conn, _, err := dialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("WebSocket dial failed: %v", err)
+	}
+	defer conn.Close()
+
+	// Send a message
+	msg := "hello"
+	if err := conn.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
+		t.Fatalf("WriteMessage failed: %v", err)
+	}
+
+	// Read the echo
+	_, p, err := conn.ReadMessage()
+	if err != nil {
+		t.Fatalf("ReadMessage failed: %v", err)
+	}
+
+	want := fmt.Sprintf("Echo: %s", msg)
+	if string(p) != want {
+		t.Errorf("Expected %q, got %q", want, p)
+	}
+}

go.mod

@@ -6,8 +6,10 @@ toolchain go1.23.11
 
 require (
 	github.com/gorilla/mux v1.8.1
+	github.com/gorilla/websocket v1.5.3
 	github.com/pkg/errors v0.8.1
 	github.com/quic-go/quic-go v0.52.0
+	golang.org/x/net v0.28.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -20,7 +22,6 @@ require (
 	go.uber.org/mock v0.5.0 // indirect
 	golang.org/x/crypto v0.26.0 // indirect
 	golang.org/x/mod v0.18.0 // indirect
-	golang.org/x/net v0.28.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.23.0 // indirect
 	golang.org/x/text v0.17.0 // indirect

go.sum

@@ -56,6 +56,8 @@ github.com/googleapis/gax-go/v2 v2.0.3/go.mod h1:LLvjysVCY1JZeum8Z6l8qUty8fiNwE0
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
 github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/grpc-gateway v1.5.0/go.mod h1:RSKVYQBd5MCa4OVpNdGskqpgL2+G+NZTnrVHpWWfpdw=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=

internal/constants/http.go

@@ -30,4 +30,8 @@ const (
 	HttpRequestURLParams               ContextKeys = "request_url_params"
 	HttpRequestPathValues              ContextKeys = "request_path_values"
 	RateLimitCustomKey                 ContextKeys = "rate_limit_custom_key"
+
+	// websocket specific context keys
+	WebsocketRequestChannel  ContextKeys = "websocket_request_channel"
+	WebsocketResponseChannel ContextKeys = "websocket_response_channel"
 )

internal/http/encoder.go

@@ -11,16 +11,9 @@ func DefaultHttpEncode(ctx context.Context, response interface{}) (headers map[s
 		"Content-Type": {"application/json; charset=utf-8"},
 	}
 
-	switch v := response.(type) {
-	case string:
-		body = []byte(v)
-	case []byte:
-		body = v
-	default:
-		body, err = json.Marshal(v)
-		if err != nil {
-			return headers, nil, err
-		}
+	body, err = GetDefaultSerialization(response)
+	if err != nil {
+		return headers, body, err
 	}
 
 	return headers, body, nil
@@ -33,3 +26,19 @@ func DefaultHttpDecode(ctx context.Context, r *http.Request) (outgoingRequest in
 
 	return outgoingRequest, nil
 }
+
+func GetDefaultSerialization(req interface{}) (body []byte, err error) {
+	switch v := req.(type) {
+	case string:
+		body = []byte(v)
+	case []byte:
+		body = v
+	default:
+		body, err = json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return body, nil
+}

internal/http/headers.go

@@ -1,19 +1,22 @@
 package http
 
-import "context"
+import (
+	"context"
+	"net/http"
+)
 
-func PopulateDefaultServerHeaders(ctx context.Context, headers map[string][]string) map[string][]string {
+func PopulateDefaultServerHeaders(ctx context.Context, r *http.Request, headers map[string][]string) map[string][]string {
 	if headers == nil {
 		headers = make(map[string][]string)
 	}
 
 	headers["X-Server"] = []string{"crazyhttp"}
-	headers["Access-Control-Allow-Origin"] = []string{"*"}
+	// relay the origin back since we check for allowed origins, earlier
+	headers["Access-Control-Allow-Origin"] = []string{r.Header.Get("Origin")}
 	headers["Access-Control-Allow-Methods"] = []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"}
 	headers["Access-Control-Allow-Headers"] = []string{"Content-Type", "Authorization"}
 	headers["Access-Control-Allow-Credentials"] = []string{"true"}
 	headers["Access-Control-Max-Age"] = []string{"86400"}
-	headers["Content-Type"] = []string{"text/event-stream"}
 
 	return headers
 }

internal/http/origin.go

@@ -0,0 +1,33 @@
+package http
+
+import (
+	"regexp"
+	"strings"
+)
+
+func IsOriginAllowed(origin string, patterns []string) bool {
+	for _, pattern := range patterns {
+		switch {
+		case len(pattern) > 2 && pattern[0] == '/' && pattern[len(pattern)-1] == '/':
+			// Treat as raw regex (/^https:\/\/foo\.com$/)
+			if matched, _ := regexp.MatchString(pattern[1:len(pattern)-1], origin); matched {
+				return true
+			}
+
+		case strings.Contains(pattern, "*"):
+			// Convert glob-style wildcard (*) to regex
+			re := "^" + regexp.QuoteMeta(pattern)
+			re = strings.ReplaceAll(re, `\*`, ".*") + "$"
+			if matched, _ := regexp.MatchString(re, origin); matched {
+				return true
+			}
+
+		default:
+			// Exact match
+			if origin == pattern {
+				return true
+			}
+		}
+	}
+	return false
+}

pkg/server/handler.go

@@ -31,6 +31,12 @@ func httpDefaultHandler(
 		return
 	}
 
+	if len(m.options.AllowedOrigins) > 0 && !ashttp.IsOriginAllowed(r.Header.Get("Origin"), m.options.AllowedOrigins) {
+		w.WriteHeader(http.StatusForbidden)
+		slog.ErrorContext(ctx, "origin not allowed", "origin", r.Header.Get("Origin"))
+		return
+	}
+
 	if decoder != nil {
 		request, err = decoder(ctx, r)
 		if err != nil {
@@ -84,7 +90,7 @@ func httpDefaultHandler(
 		}
 	}
 
-	headers = ashttp.PopulateDefaultServerHeaders(ctx, headers)
+	headers = ashttp.PopulateDefaultServerHeaders(ctx, r, headers)
 
 	for key, value := range headers {
 		w.Header().Del(key)

pkg/server/interface.go

@@ -35,6 +35,9 @@ type HttpServer interface {
 	OPTIONS(string) Method
 	CONNECT(string) Method
 	TRACE(string) Method
+
+	// Websocket
+	WebSocket(string) WebSocket
 }
 
 func NewHttpServer(ctx context.Context) HttpServer {

pkg/server/server.go

@@ -124,6 +124,10 @@ func (s *server) TRACE(url string) Method {
 	return NewMethod(constants.HttpMethodTrace, url, s)
 }
 
+func (s *server) WebSocket(url string) WebSocket {
+	return NewWebsocket(url, s)
+}
+
 // serve the HTTP request, and provide a response
 func (h *rootHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	defer func() {