feat: stream container logs to client with proper formatting

strangelove-ventures · Oct 12, 2024 · dd3681f · dd3681f
1 parent fb1be3e
commit dd3681f
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 55 deletions.
diff --git a/local-interchain/interchain/handlers/container_log_stream.go b/local-interchain/interchain/handlers/container_log_stream.go
@@ -4,6 +4,10 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"regexp"
+	"strconv"
+	"strings"
+	"unicode"
 
 	dockertypes "github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/filters"
@@ -31,7 +35,13 @@ func NewContainerSteam(ctx context.Context, logger *zap.Logger, cli *dockerclien
 }
 
 func (cs *ContainerStream) StreamContainer(w http.ResponseWriter, r *http.Request) {
-	containerID := r.URL.Query().Get("container") // TODO: get from chain ID as well? (map chain ID to container ID somehow)
+	// ensure ?auth_key=<authKey> is provided
+	if cs.authKey != "" && r.URL.Query().Get("auth_key") != cs.authKey {
+		http.Error(w, "Unauthorized, incorrect or no ?auth_key= provided", http.StatusUnauthorized)
+		return
+	}
+
+	containerID := r.URL.Query().Get("id") // TODO: get from chain ID as well? (map chain ID to container ID somehow)
 	if containerID == "" {
 		// returns containers only for this testnet. other containers are not shown on this endpoint
 		c, err := cs.cli.ContainerList(cs.ctx, dockertypes.ContainerListOptions{
@@ -56,30 +66,28 @@ func (cs *ContainerStream) StreamContainer(w http.ResponseWriter, r *http.Reques
 		return
 	}
 
+	// http://127.0.0.1:8080/container_logs?id=<ID>&colored=true
+	isColored := strings.HasPrefix(strings.ToLower(r.URL.Query().Get("colored")), "t")
+	tailLines := tailLinesParam(r.URL.Query().Get("lines"))
+
 	rr, err := cs.cli.ContainerLogs(cs.ctx, containerID, dockertypes.ContainerLogsOptions{
 		ShowStdout: true,
 		ShowStderr: true,
 		Follow:     true,
-		Tail:       "100",
-		Details:    false,
+		Details:    true,
+		Tail:       strconv.FormatUint(tailLines, 10),
 	})
 	if err != nil {
 		http.Error(w, "Unable to get container logs", http.StatusInternalServerError)
 		return
 	}
 	defer rr.Close()
 
-	// // Set headers to keep the connection open for SSE (Server-Sent Events)
+	// Set headers to keep the connection open for SSE (Server-Sent Events)
 	w.Header().Set("Content-Type", "text/event-stream")
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Connection", "keep-alive")
 
-	// ensure ?auth_key=<authKey> is provided
-	// if ls.authKey != "" && r.URL.Query().Get("auth_key") != ls.authKey {
-	// 	http.Error(w, "Unauthorized, incorrect or no ?auth_key= provided", http.StatusUnauthorized)
-	// 	return
-	// }
-
 	// Flush ensures data is sent to the client immediately
 	flusher, ok := w.(http.Flusher)
 	if !ok {
@@ -93,51 +101,52 @@ func (cs *ContainerStream) StreamContainer(w http.ResponseWriter, r *http.Reques
 		if err != nil {
 			break
 		}
-		// w.Write(buf[:n])
-		fmt.Fprintf(w, "%s", buf[:n])
-		flusher.Flush() // TODO: ?
+
+		text := string(buf[:n])
+		if !isColored {
+			text, err = removeAnsiColorCodesFromText(string(buf[:n]))
+			if err != nil {
+				http.Error(w, "Unable to remove ANSI color codes", http.StatusInternalServerError)
+				return
+			}
+		}
+
+		fmt.Fprint(w, cleanSpecialChars(text))
+		flusher.Flush()
+	}
+}
+
+func tailLinesParam(tailInput string) uint64 {
+	if tailInput == "" {
+		return defaultTailLines
 	}
 
-	// for {
-	// 	select {
-	// 	// In case client closes the connection, break out of loop
-	// 	case <-r.Context().Done():
-	// 		return
-	// 	default:
-	// 		// Try to read a line
-	// 		line, err := reader.ReadString('\n')
-	// 		if err == nil {
-	// 			// Send the log line to the client
-	// 			fmt.Fprintf(w, "%s\n", line)
-	// 			flusher.Flush() // Send to client immediately
-	// 		} else {
-	// 			// If no new log is available, wait for a short period before retrying
-	// 			time.Sleep(100 * time.Millisecond)
-	// 		}
-	// 	}
-	// }
+	tailLines, err := strconv.ParseUint(tailInput, 10, 64)
+	if err != nil {
+		return defaultTailLines
+	}
+
+	return tailLines
+}
+
+func removeAnsiColorCodesFromText(text string) (string, error) {
+	r, err := regexp.Compile("\x1b\\[[0-9;]*m")
+	if err != nil {
+		return "", err
+	}
+
+	return r.ReplaceAllString(text, ""), nil
 }
 
-// func (ls *ContainerStream) TailLogs(w http.ResponseWriter, r *http.Request) {
-// 	// ensure ?auth_key=<authKey> is provided
-// 	if ls.authKey != "" && r.URL.Query().Get("auth_key") != ls.authKey {
-// 		http.Error(w, "Unauthorized, incorrect or no ?auth_key= provided", http.StatusUnauthorized)
-// 		return
-// 	}
-
-// 	var linesToTail uint64 = defaultTailLines
-// 	tailInput := r.URL.Query().Get("lines")
-// 	if tailInput != "" {
-// 		tailLines, err := strconv.ParseUint(tailInput, 10, 64)
-// 		if err != nil {
-// 			http.Error(w, "Invalid lines input", http.StatusBadRequest)
-// 			return
-// 		}
-// 		linesToTail = tailLines
-// 	}
-
-// 	logs := TailFile(ls.logger, ls.fName, linesToTail)
-// 	for _, log := range logs {
-// 		fmt.Fprintf(w, "%s\n", log)
-// 	}
-// }
+func cleanSpecialChars(text string) string {
+	return strings.Map(func(r rune) rune {
+		if r == '\n' {
+			return r
+		}
+
+		if unicode.IsPrint(r) {
+			return r
+		}
+		return -1
+	}, text)
+}
diff --git a/local-interchain/interchain/handlers/log_stream.go b/local-interchain/interchain/handlers/log_stream.go
@@ -14,7 +14,7 @@ import (
 	"go.uber.org/zap"
 )
 
-const defaultTailLines = 25
+const defaultTailLines = 50
 
 type LogStream struct {
 	fName   string