feat: add /sse and /ws endpoints

buildx.Dockerfile

@@ -8,7 +8,7 @@ RUN apk --no-cache --no-progress add git tzdata make \
 FROM scratch
 
 COPY --from=builder /usr/share/zoneinfo /usr/share/zoneinfo
-COPY mimic /
+COPY mimic .
 
 ENTRYPOINT ["/mimic"]
 EXPOSE 80

go.mod

@@ -1,3 +1,8 @@
 module github.com/sablierapp/mimic
 
 go 1.23.2
+
+require (
+	github.com/google/uuid v1.6.0
+	github.com/gorilla/websocket v1.5.3
+)

go.sum

@@ -0,0 +1,4 @@
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=

logging.go

@@ -0,0 +1,148 @@
+package main
+
+import (
+	"bufio"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"text/template"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+type incomingTmplValues struct {
+	TraceID     string
+	Method      string
+	URL         string
+	Protocol    string
+	Accept      string
+	Host        string
+	ContentType string
+}
+
+var incomingTmpl = template.Must(template.New("").Parse(`
+Incoming Request: {{ .TraceID }}
+{{ .Method }} {{ .URL }} {{ .Protocol }}
+Accept: {{ .Accept }}
+Host: {{ .Host }}
+Content-Type: {{ .ContentType }}
+`))
+
+func fromRequest(r *http.Request) *incomingTmplValues {
+	return &incomingTmplValues{
+		TraceID:     r.Header.Get("Traceparent"),
+		Method:      r.Method,
+		URL:         r.URL.String(),
+		Protocol:    r.Proto,
+		Accept:      r.Header.Get("Accept"),
+		Host:        r.Host,
+		ContentType: r.Header.Get("Content-Type"),
+	}
+}
+
+func logRequest(r *http.Request) {
+	err := incomingTmpl.Execute(os.Stdout, fromRequest(r))
+	if err != nil {
+		log.Printf("Error executing template: %v", err)
+		return
+	}
+}
+
+var _ http.ResponseWriter = (*ResponseWriterWrapper)(nil)
+var _ http.Flusher = (*ResponseWriterWrapper)(nil)
+var _ http.Hijacker = (*ResponseWriterWrapper)(nil)
+
+type ResponseWriterWrapper struct {
+	w          *http.ResponseWriter
+	f          *http.Flusher
+	h          *http.Hijacker
+	statusCode *int
+}
+
+func (rww ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return (*rww.h).Hijack()
+}
+
+func (rww ResponseWriterWrapper) Flush() {
+	(*rww.f).Flush()
+}
+
+func (rww ResponseWriterWrapper) Header() http.Header {
+	return (*rww.w).Header()
+}
+
+func (rww ResponseWriterWrapper) Write(bytes []byte) (int, error) {
+	return (*rww.w).Write(bytes)
+}
+
+// WriteHeader function overwrites the http.ResponseWriter WriteHeader() function
+func (rww ResponseWriterWrapper) WriteHeader(statusCode int) {
+	*rww.statusCode = statusCode
+	(*rww.w).WriteHeader(statusCode)
+}
+
+// NewResponseWriterWrapper static function creates a wrapper for the http.ResponseWriter
+func NewResponseWriterWrapper(w http.ResponseWriter) ResponseWriterWrapper {
+	var statusCode int = 200
+	// Every request should implement flusher
+	flusher, _ := w.(http.Flusher)
+	// Every request should implement hijacker
+	hijacker, _ := w.(http.Hijacker)
+	return ResponseWriterWrapper{
+		w:          &w,
+		h:          &hijacker,
+		f:          &flusher,
+		statusCode: &statusCode,
+	}
+}
+
+type outgoingTmplValues struct {
+	TraceID     string
+	Duration    string
+	Protocol    string
+	StatusCode  int
+	ContentType string
+}
+
+var outgoingTmpl = template.Must(template.New("").Parse(`
+Outgoing Response: {{ .TraceID }}
+Duration: {{ .Duration }}
+{{ .Protocol }} {{ .StatusCode }}
+Content-Type: {{ .ContentType }}
+`))
+
+func fromResponse(w *ResponseWriterWrapper, r *http.Request, duration time.Duration) *outgoingTmplValues {
+	return &outgoingTmplValues{
+		TraceID:     (*w.w).Header().Get("Traceparent"),
+		Duration:    duration.String(), // To seconds
+		Protocol:    r.Proto,
+		StatusCode:  *w.statusCode,
+		ContentType: (*w.w).Header().Get("Content-Type"),
+	}
+}
+
+func logResponse(w *ResponseWriterWrapper, r *http.Request, duration time.Duration) {
+	err := outgoingTmpl.Execute(os.Stdout, fromResponse(w, r, duration))
+	if err != nil {
+		log.Printf("Error executing template: %v", err)
+		return
+	}
+}
+
+func HTTPLogging(next http.HandlerFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Generate a unique trace ID for the request
+		traceID := uuid.New().String()
+
+		// TODO: Actually add open-telemetry compatibility
+		r.Header.Set("Traceparent", traceID)
+		logRequest(r)
+		rww := NewResponseWriterWrapper(w)
+		start := time.Now()
+		next(rww, r)
+		duration := time.Since(start)
+		logResponse(&rww, r, duration)
+	})
+}

main.go

@@ -58,8 +58,10 @@ func run() int {
 
 func server() {
 	mux := http.NewServeMux()
-	mux.Handle("/health", withCLF(health))
-	mux.Handle("/", withCLF(hello))
+	mux.Handle("/health", HTTPLogging(health))
+	mux.Handle("/", HTTPLogging(hello))
+	mux.Handle("/sse", HTTPLogging(sseHandler))
+	mux.Handle("/ws", HTTPLogging(wsHandler))
 
 	log.Printf("Starting up on port %s (started in %.0f seconds)", *port, time.Since(startingTime).Seconds())
 
@@ -72,15 +74,6 @@ func server() {
 	log.Fatal(http.ListenAndServe(":"+*port, mux))
 }
 
-func withCLF(next http.HandlerFunc) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		next(w, r)
-
-		// <remote_IP_address> - [<timestamp>] "<request_method> <request_path> <request_protocol>" -
-		log.Printf("%s - - [%s] \"%s %s %s\" - -", r.RemoteAddr, time.Now().Format("02/Jan/2006:15:04:05 -0700"), r.Method, r.URL.Path, r.Proto)
-	})
-}
-
 func hello(rw http.ResponseWriter, _ *http.Request) {
 	rw.Header().Set("Content-Type", "text/plain; charset=utf-8")
 	_, err := rw.Write([]byte("Mimic says hello!"))

requests.http

@@ -0,0 +1,19 @@
+GET http://localhost:80/
+
+###
+GET http://localhost:80/health
+
+###
+WEBSOCKET ws://localhost:80/ws
+
+=== wait-for-server
+Hi Server!
+=== wait-for-server
+Are you repeating what I'm saying?
+=== wait-for-server
+Ok bye.
+
+
+###
+GET http://localhost:80/sse
+Accept: text/event-stream

sse.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"time"
+)
+
+func sseHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+
+	// Flush the headers
+	flusher, ok := w.(http.Flusher)
+	if !ok {
+		http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
+		return
+	}
+
+	// Start sending periodic events to the client every second
+	ticker := time.NewTicker(1 * time.Second)
+	defer ticker.Stop()
+
+	// Send events every second
+	// Send events every second
+	for {
+		select {
+		case <-ticker.C:
+			// Send an event with the current timestamp
+			fmt.Fprintf(w, "data: Current time: %s\n\n", time.Now().Format(time.RFC3339))
+			flusher.Flush()
+		case <-r.Context().Done():
+			log.Printf("Request cancelled: %v", r.Context().Err())
+			return
+		}
+	}
+}

websoket.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"fmt"
+	"github.com/gorilla/websocket"
+	"log"
+	"net/http"
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	CheckOrigin:     func(_ *http.Request) bool { return true },
+}
+
+func reader(conn *websocket.Conn) {
+	for {
+		// read in a message
+		messageType, p, err := conn.ReadMessage()
+		if err != nil {
+			log.Println(err)
+			return
+		}
+
+		response := fmt.Sprintf("You said: %s", string(p))
+		err = conn.WriteMessage(1, []byte(response))
+		if err != nil {
+			log.Println(err)
+		}
+		if err := conn.WriteMessage(messageType, p); err != nil {
+			log.Println(err)
+			return
+		}
+	}
+}
+
+func wsHandler(w http.ResponseWriter, r *http.Request) {
+	// upgrade this connection to a WebSocket
+	// connection
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer ws.Close()
+
+	log.Println("Client Connected")
+	err = ws.WriteMessage(1, []byte("Hi Client!"))
+	if err != nil {
+		log.Println(err)
+	}
+	// listen indefinitely for new messages coming
+	// through on our WebSocket connection
+	reader(ws)
+}