gochat/backend/internal/middleware/xss_protection.go

package middleware

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/microcosm-cc/bluemonday"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// XSSProtectionConfig holds configuration for the XSS protection middleware.
type XSSProtectionConfig struct {
	// HTMLPolicy is the bluemonday policy used for HTML sanitization.
	// If nil, a strict default policy is used that strips all HTML tags.
	HTMLPolicy *bluemonday.Policy
	// SanitizeJSONResponse enables HTML escaping in JSON responses.
	SanitizeJSONResponse bool
	// SanitizeInputFields enables sanitization of form/query parameters.
	SanitizeInputFields bool
	// InputFieldsToSanitize lists specific field names to sanitize from input.
	// If empty, all common text fields (content, message, description, name, title, body) are sanitized.
	InputFieldsToSanitize []string
}

// DefaultXSSProtectionConfig returns a strict default configuration.
func DefaultXSSProtectionConfig() *XSSProtectionConfig {
	policy := bluemonday.StrictPolicy() // Strip all HTML tags — strict sanitization

	return &XSSProtectionConfig{
		HTMLPolicy:            policy,
		SanitizeJSONResponse:  true,
		SanitizeInputFields:   true,
		InputFieldsToSanitize: []string{
			"content", "message", "description", "name",
			"title", "body", "subject", "comment", "note",
			"bio", "display_name", "email",
		},
	}
}

// UGCPolicy returns a policy that allows safe user-generated content HTML.
// This permits basic formatting tags (b, i, em, strong, a with href, p, br, ul, ol, li)
// while stripping dangerous elements (script, iframe, object, etc.).
func UGCPolicy() *bluemonday.Policy {
	p := bluemonday.UGCPolicy()
	// Restrict links to safe protocols only
	p.AllowStandardURLs()
	return p
}

// SanitizeHTML cleans HTML content using bluemonday with the given policy.
// If policy is nil, the strict policy (strip all tags) is used.
func SanitizeHTML(input string, policy *bluemonday.Policy) string {
	if policy == nil {
		policy = bluemonday.StrictPolicy()
	}
	return policy.Sanitize(input)
}

// EscapeJSONHTML safely marshals data to JSON with HTML-escaped strings.
// This prevents XSS when JSON responses are embedded in HTML pages.
// Go's encoding/json already escapes <, >, and & by default, but this
// function provides an extra layer of safety by also escaping single quotes
// and ensuring proper handling of unsafe unicode sequences.
func EscapeJSONHTML(data interface{}) ([]byte, error) {
	raw, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}

	// Go's json.Marshal already escapes &, <, >, and \u2028/\u2029.
	// We additionally escape single quotes for inline script safety.
	buf := bytes.ReplaceAll(raw, []byte("'"), []byte("\\u0027"))

	return buf, nil
}

// sanitizeInputValue sanitizes a string value for XSS prevention.
// This removes HTML tags and escapes dangerous characters.
func sanitizeInputValue(value string, policy *bluemonday.Policy) string {
	if value == "" {
		return value
	}

	// First strip any HTML tags using bluemonday
	cleaned := SanitizeHTML(value, policy)

	// Additionally remove null bytes and control characters that could be
	// used for injection attacks
	var result []rune
	for _, r := range cleaned {
		if r < 32 && r != '\t' && r != '\n' && r != '\r' {
			continue // skip control characters
		}
		if r == 0 {
			continue // skip null bytes
		}
		result = append(result, r)
	}

	return string(result)
}

// sanitizeRecursive walks through interface{} values and sanitizes string fields
// found in maps or slices. This is used for JSON body sanitization.
func sanitizeRecursive(data interface{}, policy *bluemonday.Policy, fieldsToSanitize map[string]bool) interface{} {
	switch v := data.(type) {
	case map[string]interface{}:
		result := make(map[string]interface{}, len(v))
		for key, val := range v {
			if fieldsToSanitize[key] {
				if strVal, ok := val.(string); ok {
					result[key] = sanitizeInputValue(strVal, policy)
				} else {
					result[key] = sanitizeRecursive(val, policy, fieldsToSanitize)
				}
			} else {
				result[key] = sanitizeRecursive(val, policy, fieldsToSanitize)
			}
		}
		return result
	case []interface{}:
		result := make([]interface{}, len(v))
		for i, val := range v {
			result[i] = sanitizeRecursive(val, policy, fieldsToSanitize)
		}
		return result
	case string:
		return sanitizeInputValue(v, policy)
	default:
		return v
	}
}

// shouldSanitizeField determines if a field name should be sanitized.
func shouldSanitizeField(fieldName string, config *XSSProtectionConfig) bool {
	fieldLower := strings.ToLower(fieldName)
	for _, f := range config.InputFieldsToSanitize {
		if strings.ToLower(f) == fieldLower {
			return true
		}
	}
	return false
}

// XSSProtectionMiddleware provides comprehensive XSS protection:
//   - HTML sanitization of input fields using bluemonday
//   - JSON response HTML escaping for safe embedding in HTML pages
//   - Input cleansing of form/query/body parameters
//   - Content-Type header enforcement
func XSSProtectionMiddleware(cfg *XSSProtectionConfig) gin.HandlerFunc {
	if cfg == nil {
		cfg = DefaultXSSProtectionConfig()
	}
	if cfg.HTMLPolicy == nil {
		cfg.HTMLPolicy = bluemonday.StrictPolicy()
	}

	// Build set of fields to sanitize for quick lookup
	fieldsSet := make(map[string]bool, len(cfg.InputFieldsToSanitize))
	for _, f := range cfg.InputFieldsToSanitize {
		fieldsSet[f] = true
	}

	return func(c *gin.Context) {
		// 1. Set XSS prevention security headers
		c.Header("X-Content-Type-Options", "nosniff")
		c.Header("X-XSS-Protection", "1; mode=block")
		c.Header("Referrer-Policy", "strict-origin-when-cross-origin")
		c.Header("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'")

		// 2. Sanitize input query/form parameters
		if cfg.SanitizeInputFields {
			// Sanitize query parameters
			for key, values := range c.Request.URL.Query() {
				if shouldSanitizeField(key, cfg) {
					sanitizedValues := make([]string, len(values))
					for i, v := range values {
						sanitizedValues[i] = sanitizeInputValue(v, cfg.HTMLPolicy)
					}
					// Replace query values by updating the URL
					q := c.Request.URL.Query()
					q[key] = sanitizedValues
					c.Request.URL.RawQuery = q.Encode()
				}
			}

			// Sanitize form POST parameters
			if c.Request.Method == http.MethodPost || c.Request.Method == http.MethodPut || c.Request.Method == http.MethodPatch {
				if err := c.Request.ParseForm(); err == nil {
					for key, values := range c.Request.PostForm {
						if shouldSanitizeField(key, cfg) {
							sanitizedValues := make([]string, len(values))
							for i, v := range values {
								sanitizedValues[i] = sanitizeInputValue(v, cfg.HTMLPolicy)
							}
							c.Request.PostForm[key] = sanitizedValues
						}
					}
				}
			}

			// Sanitize JSON body
			if c.Request.Method == http.MethodPost || c.Request.Method == http.MethodPut || c.Request.Method == http.MethodPatch {
				contentType := c.GetHeader("Content-Type")
				if contentType != "" && contentType == "application/json" {
					if c.Request.Body != nil {
						bodyBytes, err := io.ReadAll(c.Request.Body)
						if err == nil && len(bodyBytes) > 0 {
							var bodyData interface{}
							if err := json.Unmarshal(bodyBytes, &bodyData); err == nil {
								sanitizedData := sanitizeRecursive(bodyData, cfg.HTMLPolicy, fieldsSet)
								sanitizedBytes, err := json.Marshal(sanitizedData)
								if err == nil {
									c.Request.Body = io.NopCloser(bytes.NewReader(sanitizedBytes))
									// Update Content-Length to reflect potentially changed body
									c.Request.ContentLength = int64(len(sanitizedBytes))
								}
							} else {
								// Restore original body if JSON parse fails
								c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
							}
						}
					}
				}
			}
		}

		// 3. JSON response HTML escaping — intercept c.JSON calls
		if cfg.SanitizeJSONResponse {
			// We replace the writer to ensure HTML-safe JSON output
			originalWriter := c.Writer
			xssWriter := &xssResponseWriter{
				ResponseWriter: originalWriter,
				policy:         cfg.HTMLPolicy,
			}
			c.Writer = xssWriter
		}

		c.Next()

		// Restore original writer after handler completes
		if cfg.SanitizeJSONResponse {
			c.Writer = c.Writer.(*xssResponseWriter).ResponseWriter
		}
	}
}

// xssResponseWriter wraps gin.ResponseWriter to provide XSS-safe JSON output.
type xssResponseWriter struct {
	gin.ResponseWriter
	policy *bluemonday.Policy
}

// WriteJSON overrides the default JSON writing to apply HTML escaping.
func (w *xssResponseWriter) WriteJSON(data interface{}) error {
	escaped, err := EscapeJSONHTML(data)
	if err != nil {
		applogger.L().Errorf("XSS protection: failed to escape JSON response: %v", err)
		return err
	}
	_, writeErr := w.ResponseWriter.Write(escaped)
	return writeErr
}

// Write ensures that any written content is also XSS-safe for JSON content types.
func (w *xssResponseWriter) Write(data []byte) (int, error) {
	// Only escape if the content type is JSON
	contentType := w.Header().Get("Content-Type")
	if strings.HasPrefix(contentType, "application/json") {
		// The data is already JSON-encoded, just escape single quotes for safety
		escaped := bytes.ReplaceAll(data, []byte("'"), []byte("\\u0027"))
		return w.ResponseWriter.Write(escaped)
	}
	return w.ResponseWriter.Write(data)
}