feat: c.IP add use-proxy-header arg

.gitignore

@@ -1,3 +1,4 @@
 .idea
 .vscode
 .DS_Store
+xtest

app.go

@@ -19,8 +19,6 @@ var (
 
 	regSafePrefix         = regexp.MustCompile("[^a-zA-Z0-9/-]+")
 	regRemoveRepeatedChar = regexp.MustCompile("/{2,}")
-
-	mimePlain = []string{"text/plain"}
 )
 
 type App struct {
@@ -176,7 +174,6 @@ func (a *App) handleHTTPRequest(c *Ctx) {
 				serveError(c, errorHandler)
 			}
 
-			c.writermem.WriteHeaderNow()
 			return
 		}
 		if httpMethod != http.MethodConnect && rPath != "/" {
@@ -205,17 +202,19 @@ func (a *App) handleHTTPRequest(c *Ctx) {
 		}
 
 		if len(allowed) > 0 {
-			c.handlers = a.combineHandlers()
+			c.handlers = a.combineHandlers(a.config.MethodNotAllowedHandler)
 
-			serveError(c, a.config.MethodNotAllowedHandler)
+			_ = c.Next()
 
 			return
 		}
 	}
 
-	c.handlers = a.combineHandlers()
+	c.handlers = a.combineHandlers(a.config.NotFoundHandler)
 
-	serveError(c, a.config.NotFoundHandler)
+	_ = c.Next()
+
+	return
 }
 
 func errorHandler(c *Ctx) error {
@@ -274,6 +273,6 @@ func redirectRequest(c *Ctx) {
 
 	//debugPrint("redirecting request %d: %s --> %s", code, rPath, rURL)
 
-	http.Redirect(c.writer, req, rURL, code)
+	http.Redirect(c.Writer, req, rURL, code)
 	c.writermem.WriteHeaderNow()
 }

ctx.go

@@ -3,7 +3,9 @@ package nf
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"github.com/loveuer/nf/internal/sse"
 	"io"
 	"mime/multipart"
 	"net"
@@ -12,10 +14,14 @@ import (
 	"sync"
 )
 
+var (
+	forwardHeaders = []string{"CF-Connecting-IP", "X-Forwarded-For", "X-Real-Ip"}
+)
+
 type Ctx struct {
 	lock       sync.Mutex
 	writermem  responseWriter
-	writer     http.ResponseWriter
+	Writer     ResponseWriter
 	Request    *http.Request
 	path       string
 	method     string
@@ -37,7 +43,6 @@ func newContext(app *App, writer http.ResponseWriter, request *http.Request) *Ct
 
 	ctx := &Ctx{
 		lock:       sync.Mutex{},
-		writer:     writer,
 		Request:    request,
 		path:       request.URL.Path,
 		method:     request.Method,
@@ -52,11 +57,13 @@ func newContext(app *App, writer http.ResponseWriter, request *http.Request) *Ct
 	}
 
 	ctx.writermem = responseWriter{
-		ResponseWriter: ctx.writer,
+		ResponseWriter: writer,
 		size:           -1,
 		status:         0,
 	}
 
+	ctx.Writer = &ctx.writermem
+
 	return ctx
 }
 
@@ -188,11 +195,24 @@ func (c *Ctx) Get(key string, defaultValue ...string) string {
 	return value
 }
 
-func (c *Ctx) IP() string {
+func (c *Ctx) IP(useProxyHeader ...bool) string {
 	ip, _, err := net.SplitHostPort(strings.TrimSpace(c.Request.RemoteAddr))
 	if err != nil {
 		return ""
 	}
+
+	if len(useProxyHeader) > 0 && useProxyHeader[0] {
+		for _, h := range forwardHeaders {
+			for _, rip := range strings.Split(c.Request.Header.Get(h), ",") {
+				realIP := net.ParseIP(strings.Replace(rip, " ", "", -1))
+				if check := net.ParseIP(realIP.String()); check != nil {
+					ip = realIP.String()
+					break
+				}
+			}
+		}
+	}
+
 	return ip
 }
 
@@ -253,8 +273,8 @@ func (c *Ctx) Status(code int) *Ctx {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	c.StatusCode = code
-	c.writermem.status = code
+	c.writermem.WriteHeader(code)
+	c.StatusCode = c.writermem.status
 
 	return c
 }
@@ -281,7 +301,7 @@ func (c *Ctx) SendString(data string) error {
 
 func (c *Ctx) Writef(format string, values ...interface{}) (int, error) {
 	c.SetHeader("Content-Type", "text/plain")
-	return c.writer.Write([]byte(fmt.Sprintf(format, values...)))
+	return c.Write([]byte(fmt.Sprintf(format, values...)))
 }
 
 func (c *Ctx) JSON(data interface{}) error {
@@ -296,13 +316,26 @@ func (c *Ctx) JSON(data interface{}) error {
 	return nil
 }
 
-func (c *Ctx) RawWriter() http.ResponseWriter {
-	return c.writer
+func (c *Ctx) SSEvent(event string, data interface{}) error {
+	c.Set("Content-Type", "text/event-stream")
+	c.Set("Cache-Control", "no-cache")
+	c.Set("Transfer-Encoding", "chunked")
+
+	return sse.Encode(c.Writer, sse.Event{Event: event, Data: data})
+}
+
+func (c *Ctx) Flush() error {
+	if f, ok := c.Writer.(http.Flusher); ok {
+		f.Flush()
+		return nil
+	}
+
+	return errors.New("http.Flusher is not implemented")
 }
 
 func (c *Ctx) HTML(html string) error {
 	c.SetHeader("Content-Type", "text/html")
-	_, err := c.writer.Write([]byte(html))
+	_, err := c.Write([]byte(html))
 	return err
 }
 

go.mod

@@ -1,3 +1,14 @@
 module github.com/loveuer/nf
 
 go 1.20
+
+require (
+	github.com/fatih/color v1.17.0
+	github.com/google/uuid v1.6.0
+)
+
+require (
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	golang.org/x/sys v0.18.0 // indirect
+)

go.sum

@@ -0,0 +1,13 @@
+github.com/fatih/color v1.17.0 h1:GlRw1BRJxkpqUCBKzKOw098ed57fEsKeNjpTe3cSjK4=
+github.com/fatih/color v1.17.0/go.mod h1:YZ7TlrGPkiz6ku9fK3TLD/pl3CpsiFyu8N92HLgmosI=
+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/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
+golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=

internal/sse/sse-encoder.go

@@ -0,0 +1,106 @@
+package sse
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"reflect"
+	"strconv"
+	"strings"
+)
+
+// Server-Sent Events
+// W3C Working Draft 29 October 2009
+// http://www.w3.org/TR/2009/WD-eventsource-20091029/
+
+const ContentType = "text/event-stream"
+
+var contentType = []string{ContentType}
+var noCache = []string{"no-cache"}
+
+var fieldReplacer = strings.NewReplacer(
+	"\n", "\\n",
+	"\r", "\\r")
+
+var dataReplacer = strings.NewReplacer(
+	"\n", "\ndata:",
+	"\r", "\\r")
+
+type Event struct {
+	Event string
+	Id    string
+	Retry uint
+	Data  interface{}
+}
+
+func Encode(writer io.Writer, event Event) error {
+	w := checkWriter(writer)
+	writeId(w, event.Id)
+	writeEvent(w, event.Event)
+	writeRetry(w, event.Retry)
+	return writeData(w, event.Data)
+}
+
+func writeId(w stringWriter, id string) {
+	if len(id) > 0 {
+		w.WriteString("id:")
+		fieldReplacer.WriteString(w, id)
+		w.WriteString("\n")
+	}
+}
+
+func writeEvent(w stringWriter, event string) {
+	if len(event) > 0 {
+		w.WriteString("event:")
+		fieldReplacer.WriteString(w, event)
+		w.WriteString("\n")
+	}
+}
+
+func writeRetry(w stringWriter, retry uint) {
+	if retry > 0 {
+		w.WriteString("retry:")
+		w.WriteString(strconv.FormatUint(uint64(retry), 10))
+		w.WriteString("\n")
+	}
+}
+
+func writeData(w stringWriter, data interface{}) error {
+	w.WriteString("data:")
+	switch kindOfData(data) {
+	case reflect.Struct, reflect.Slice, reflect.Map:
+		err := json.NewEncoder(w).Encode(data)
+		if err != nil {
+			return err
+		}
+		w.WriteString("\n")
+	default:
+		dataReplacer.WriteString(w, fmt.Sprint(data))
+		w.WriteString("\n\n")
+	}
+	return nil
+}
+
+func (r Event) Render(w http.ResponseWriter) error {
+	r.WriteContentType(w)
+	return Encode(w, r)
+}
+
+func (r Event) WriteContentType(w http.ResponseWriter) {
+	header := w.Header()
+	header["Content-Type"] = contentType
+
+	if _, exist := header["Cache-Control"]; !exist {
+		header["Cache-Control"] = noCache
+	}
+}
+
+func kindOfData(data interface{}) reflect.Kind {
+	value := reflect.ValueOf(data)
+	valueType := value.Kind()
+	if valueType == reflect.Ptr {
+		valueType = value.Elem().Kind()
+	}
+	return valueType
+}

internal/sse/writer.go

@@ -0,0 +1,24 @@
+package sse
+
+import "io"
+
+type stringWriter interface {
+	io.Writer
+	WriteString(string) (int, error)
+}
+
+type stringWrapper struct {
+	io.Writer
+}
+
+func (w stringWrapper) WriteString(str string) (int, error) {
+	return w.Writer.Write([]byte(str))
+}
+
+func checkWriter(writer io.Writer) stringWriter {
+	if w, ok := writer.(stringWriter); ok {
+		return w
+	} else {
+		return stringWrapper{writer}
+	}
+}

middleware.go

@@ -2,9 +2,11 @@ package nf
 
 import (
 	"fmt"
-	"log"
+	"github.com/google/uuid"
+	"github.com/loveuer/nf/nft/log"
 	"os"
 	"runtime/debug"
+	"strings"
 	"time"
 )
 
@@ -27,51 +29,44 @@ func NewRecover(enableStackTrace bool) HandlerFunc {
 	}
 }
 
-func NewLogger() HandlerFunc {
-	l := log.New(os.Stdout, "[NF] ", 0)
-
-	durationFormat := func(num int64) string {
-		var (
-			unit = "ns"
-		)
-
-		if num > 1000 {
-			num = num / 1000
-			unit = "µs"
-		}
-
-		if num > 1000 {
-			num = num / 1000
-			unit = "ms"
-		}
-
-		if num > 1000 {
-			num = num / 1000
-			unit = "s"
-		}
-
-		return fmt.Sprintf("%3d %2s", num, unit)
+func NewLogger(traceHeader ...string) HandlerFunc {
+	Header := "X-Trace-ID"
+	if len(traceHeader) > 0 && traceHeader[0] != "" {
+		Header = traceHeader[0]
 	}
 
 	return func(c *Ctx) error {
-		start := time.Now()
+		var (
+			now   = time.Now()
+			trace = c.Get(Header)
+			logFn func(msg string, data ...any)
+			ip    = c.IP()
+		)
+
+		if trace == "" {
+			trace = uuid.Must(uuid.NewV7()).String()
+		}
+
+		c.SetHeader(Header, trace)
+
+		traces := strings.Split(trace, "-")
+		shortTrace := traces[len(traces)-1]
 
 		err := c.Next()
+		duration := time.Since(now)
 
-		var (
-			duration = time.Now().Sub(start).Nanoseconds()
-			status   = c.StatusCode
-			path     = c.path
-			method   = c.Request.Method
-		)
+		msg := fmt.Sprintf("NF | %s | %15s | %3d | %s | %6s | %s", shortTrace, ip, c.StatusCode, HumanDuration(duration.Nanoseconds()), c.Method(), c.Path())
 
-		l.Printf("%s | %5s | %d | %s | %s",
-			start.Format("06/01/02T15:04:05"),
-			method,
-			status,
-			durationFormat(duration),
-			path,
-		)
+		switch {
+		case c.StatusCode >= 500:
+			logFn = log.Error
+		case c.StatusCode >= 400:
+			logFn = log.Warn
+		default:
+			logFn = log.Info
+		}
+
+		logFn(msg)
 
 		return err
 	}

nft/log/default.go

@@ -0,0 +1,67 @@
+package log
+
+import (
+	"fmt"
+	"os"
+	"sync"
+)
+
+var (
+	nilLogger    = func(prefix, timestamp, msg string, data ...any) {}
+	normalLogger = func(prefix, timestamp, msg string, data ...any) {
+		fmt.Printf(prefix+"| "+timestamp+" | "+msg+"\n", data...)
+	}
+
+	panicLogger = func(prefix, timestamp, msg string, data ...any) {
+		panic(fmt.Sprintf(prefix+"| "+timestamp+" | "+msg+"\n", data...))
+	}
+
+	fatalLogger = func(prefix, timestamp, msg string, data ...any) {
+		fmt.Printf(prefix+"| "+timestamp+" | "+msg+"\n", data...)
+		os.Exit(1)
+	}
+
+	defaultLogger = &logger{
+		Mutex:      sync.Mutex{},
+		timeFormat: "2006-01-02T15:04:05",
+		writer:     os.Stdout,
+		level:      LogLevelInfo,
+		debug:      nilLogger,
+		info:       normalLogger,
+		warn:       normalLogger,
+		error:      normalLogger,
+		panic:      panicLogger,
+		fatal:      fatalLogger,
+	}
+)
+
+func SetTimeFormat(format string) {
+	defaultLogger.SetTimeFormat(format)
+}
+
+func SetLogLevel(level LogLevel) {
+	defaultLogger.SetLogLevel(level)
+}
+
+func Debug(msg string, data ...any) {
+	defaultLogger.Debug(msg, data...)
+}
+func Info(msg string, data ...any) {
+	defaultLogger.Info(msg, data...)
+}
+
+func Warn(msg string, data ...any) {
+	defaultLogger.Warn(msg, data...)
+}
+
+func Error(msg string, data ...any) {
+	defaultLogger.Error(msg, data...)
+}
+
+func Panic(msg string, data ...any) {
+	defaultLogger.Panic(msg, data...)
+}
+
+func Fatal(msg string, data ...any) {
+	defaultLogger.Fatal(msg, data...)
+}

nft/log/log.go

@@ -0,0 +1,115 @@
+package log
+
+import (
+	"github.com/fatih/color"
+	"io"
+	"sync"
+	"time"
+)
+
+type LogLevel uint32
+
+const (
+	LogLevelDebug = iota
+	LogLevelInfo
+	LogLevelWarn
+	LogLevelError
+	LogLevelPanic
+	LogLevelFatal
+)
+
+type logger struct {
+	sync.Mutex
+	timeFormat string
+	writer     io.Writer
+	level      LogLevel
+	debug      func(prefix, timestamp, msg string, data ...any)
+	info       func(prefix, timestamp, msg string, data ...any)
+	warn       func(prefix, timestamp, msg string, data ...any)
+	error      func(prefix, timestamp, msg string, data ...any)
+	panic      func(prefix, timestamp, msg string, data ...any)
+	fatal      func(prefix, timestamp, msg string, data ...any)
+}
+
+var (
+	red    = color.New(color.FgRed)
+	hired  = color.New(color.FgHiRed)
+	green  = color.New(color.FgGreen)
+	yellow = color.New(color.FgYellow)
+	white  = color.New(color.FgWhite)
+)
+
+func (l *logger) SetTimeFormat(format string) {
+	l.Lock()
+	defer l.Unlock()
+	l.timeFormat = format
+}
+
+func (l *logger) SetLogLevel(level LogLevel) {
+	l.Lock()
+	defer l.Unlock()
+
+	if level > LogLevelDebug {
+		l.debug = nilLogger
+	} else {
+		l.debug = normalLogger
+	}
+
+	if level > LogLevelInfo {
+		l.info = nilLogger
+	} else {
+		l.info = normalLogger
+	}
+
+	if level > LogLevelWarn {
+		l.warn = nilLogger
+	} else {
+		l.warn = normalLogger
+	}
+
+	if level > LogLevelError {
+		l.error = nilLogger
+	} else {
+		l.error = normalLogger
+	}
+
+	if level > LogLevelPanic {
+		l.panic = nilLogger
+	} else {
+		l.panic = panicLogger
+	}
+
+	if level > LogLevelFatal {
+		l.fatal = nilLogger
+	} else {
+		l.fatal = fatalLogger
+	}
+}
+
+func (l *logger) Debug(msg string, data ...any) {
+	l.debug(white.Sprint("Debug "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+func (l *logger) Info(msg string, data ...any) {
+	l.info(green.Sprint("Info  "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+func (l *logger) Warn(msg string, data ...any) {
+	l.warn(yellow.Sprint("Warn  "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+func (l *logger) Error(msg string, data ...any) {
+	l.error(red.Sprint("Error "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+func (l *logger) Panic(msg string, data ...any) {
+	l.panic(hired.Sprint("Panic "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+func (l *logger) Fatal(msg string, data ...any) {
+	l.fatal(hired.Sprint("Fatal "), time.Now().Format(l.timeFormat), msg, data...)
+}
+
+type WroteLogger interface {
+	Info(msg string, data ...any)
+}

nft/log/new.go

@@ -0,0 +1,21 @@
+package log
+
+import (
+	"os"
+	"sync"
+)
+
+func New() *logger {
+	return &logger{
+		Mutex:      sync.Mutex{},
+		timeFormat: "2006-01-02T15:04:05",
+		writer:     os.Stdout,
+		level:      LogLevelInfo,
+		debug:      nilLogger,
+		info:       normalLogger,
+		warn:       normalLogger,
+		error:      normalLogger,
+		panic:      panicLogger,
+		fatal:      fatalLogger,
+	}
+}

response_writer.go

@@ -3,6 +3,7 @@ package nf
 import (
 	"bufio"
 	"io"
+	"log"
 	"net"
 	"net/http"
 )
@@ -41,6 +42,7 @@ type ResponseWriter interface {
 
 type responseWriter struct {
 	http.ResponseWriter
+	written bool
 	size    int
 	status  int
 }
@@ -60,7 +62,7 @@ func (w *responseWriter) reset(writer http.ResponseWriter) {
 func (w *responseWriter) WriteHeader(code int) {
 	if code > 0 && w.status != code {
 		if w.Written() {
-			// todo: debugPrint("[WARNING] Headers were already written. Wanted to override status code %d with %d", w.status, code)
+			log.Printf("[NF] WARNING: Headers were already written. Wanted to override status code %d with %d", w.status, code)
 			return
 		}
 		w.status = code
@@ -102,7 +104,7 @@ func (w *responseWriter) Size() int {
 }
 
 func (w *responseWriter) Written() bool {
-	return w.size != noWritten
+	return w.size != noWritten || w.written
 }
 
 // Hijack implements the http.Hijacker interface.

util.go

@@ -202,3 +202,24 @@ func bufApp(buf *[]byte, s string, w int, c byte) {
 	}
 	b[w] = c
 }
+
+func HumanDuration(nano int64) string {
+	duration := float64(nano)
+	unit := "ns"
+	if duration >= 1000 {
+		duration /= 1000
+		unit = "us"
+	}
+
+	if duration >= 1000 {
+		duration /= 1000
+		unit = "ms"
+	}
+
+	if duration >= 1000 {
+		duration /= 1000
+		unit = " s"
+	}
+
+	return fmt.Sprintf("%6.2f%s", duration, unit)
+}