rano/bin/server/handler/request.go

// Copyright 2024 Patial Tech. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package handler

import (
	"context"
	"net"
	"net/http"
	"strings"

	"gitserver.in/patialtech/rano/util/uid"
)

const RequestIDKey = "RequestID"
const RequestIPKey = "RequestIP"
const RequestUserAgentKey = "RequestUA"

var defaultHeaders = []string{
	"True-Client-IP", // Cloudflare Enterprise plan
	"X-Real-IP",
	"X-Forwarded-For",
}

// Request middleware that will do the following:
//   - pull session user
//   - set requestID
//   - set ctx RealIP and client userAgent info
func Request() func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()

			// ID
			requestID := r.Header.Get("X-Request-Id")
			if requestID == "" {
				requestID = uid.NewUUID()
			}
			ctx = context.WithValue(ctx, RequestIDKey, requestID)

			// IP
			if ip := getRealIP(r, defaultHeaders); ip != "" {
				r.RemoteAddr = ip
			}
			ctx = context.WithValue(ctx, RequestIPKey, requestID)

			// User Agent
			ctx = context.WithValue(ctx, RequestUserAgentKey, r.UserAgent())

			h.ServeHTTP(w, r)
		})
	}
}

func getRealIP(r *http.Request, headers []string) string {
	for _, header := range headers {
		if ip := r.Header.Get(header); ip != "" {
			ips := strings.Split(ip, ",")
			if ips[0] == "" || net.ParseIP(ips[0]) == nil {
				continue
			}
			return ips[0]
		}
	}
	return ""
}