package httpserver

import (
	"bytes"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type loginBucket struct {
	Start time.Time
	Count int
}

func (s *Server) allowLogin(address string) bool {
	ip, _, err := net.SplitHostPort(address)
	if err != nil {
		ip = address
	}
	s.loginMu.Lock()
	defer s.loginMu.Unlock()
	if s.logins == nil {
		s.logins = map[string]loginBucket{}
	}
	now := time.Now()
	for k, v := range s.logins {
		if now.Sub(v.Start) > time.Minute {
			delete(s.logins, k)
		}
	}
	b := s.logins[ip]
	if b.Start.IsZero() {
		if len(s.logins) >= 1024 {
			return false
		}
		b.Start = now
	}
	b.Count++
	s.logins[ip] = b
	return b.Count <= 10
}

func requestProtection(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// No proxy headers are trusted. Configure the proxy to preserve the Host.
		admin := strings.HasPrefix(r.URL.Path, "/ui/") || r.URL.Path == "/login" || r.URL.Path == "/logout"
		if admin {
			w.Header().Set("Cache-Control", "no-store")
		}
		if admin && r.Method != "GET" && r.Method != "HEAD" {
			if origin := r.Header.Get("Origin"); origin != "" {
				u, err := url.Parse(origin)
				if err != nil || u.Host != r.Host || (u.Scheme != "http" && u.Scheme != "https") {
					http.Error(w, "cross-origin request rejected", 403)
					return
				}
			}
			if r.Header.Get("Sec-Fetch-Site") == "cross-site" {
				http.Error(w, "cross-site request rejected", 403)
				return
			}
		}
		if r.Body != nil {
			limit := int64(1 << 20)
			if admin {
				limit = 2 << 20
			}
			b, err := io.ReadAll(io.LimitReader(r.Body, limit+1))
			r.Body.Close()
			if err != nil {
				http.Error(w, "invalid body", 400)
				return
			}
			if int64(len(b)) > limit {
				http.Error(w, "request too large", 413)
				return
			}
			r.Body = io.NopCloser(bytes.NewReader(b))
		}
		next.ServeHTTP(w, r)
	})
}