From b6e7396ef50951976689dc03266007246e105cea Mon Sep 17 00:00:00 2001
From: kingecg <kingecg@hotmail.com>
Date: Sat, 9 Dec 2023 16:34:20 +0800
Subject: [PATCH] add handler

---
 hander/file.go  | 56 ++++++++++++++++++++++++++++++++
 hander/proxy.go | 86 +++++++++++++++++++++++++++++++++++++++++++++++++
 model/model.go  | 14 +++++---
 3 files changed, 152 insertions(+), 4 deletions(-)
 create mode 100644 hander/file.go
 create mode 100644 hander/proxy.go

diff --git a/hander/file.go b/hander/file.go
new file mode 100644
index 0000000..92ebcca
--- /dev/null
+++ b/hander/file.go
@@ -0,0 +1,56 @@
+package handler
+
+import (
+	"errors"
+	"io/fs"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+type FileHandler struct {
+	Root    string
+	Default string
+	http.FileSystem
+}
+
+func (f FileHandler) Open(name string) (http.File, error) {
+	if strings.HasPrefix(name, "../") {
+		return nil, errors.New("not permitted")
+	}
+
+	rPath := filepath.Join(f.Root, strings.TrimPrefix(name, "/"))
+	if rPath == f.Root {
+		if f.Default == "" {
+			return nil, errors.New("not permit list dir")
+		}
+		rPath = filepath.Join(rPath, f.Default)
+	}
+
+	fInfo, _, err := FileExists(rPath)
+	if err != nil {
+		return nil, err
+	}
+
+	if fInfo.IsDir() {
+		return nil, errors.New("not permit list dir")
+	}
+
+	if fInfo.Mode() == fs.ModeSymlink {
+		return nil, errors.New("not permit follow symbol link")
+	}
+	return os.Open(rPath)
+}
+
+func FileExists(name string) (os.FileInfo, bool, error) {
+	info, err := os.Stat(name)
+	if err != nil {
+		return nil, os.IsExist(err), err
+	}
+	return info, true, nil
+}
+
+func NewFileHandler(f FileHandler) http.Handler {
+	return http.FileServer(f)
+}
diff --git a/hander/proxy.go b/hander/proxy.go
new file mode 100644
index 0000000..7cfe163
--- /dev/null
+++ b/hander/proxy.go
@@ -0,0 +1,86 @@
+package handler
+
+import (
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"strconv"
+	"strings"
+
+	"git.pyer.club/kingecg/gohttpd/model"
+)
+
+type ProxyHandler struct {
+	proxy []*httputil.ReverseProxy
+	count int
+}
+
+func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	s, err := r.Cookie("s")
+	var proxyIndex int
+	if err != nil {
+		proxyIndex = p.count
+		p.count++
+		if p.count >= len(p.proxy) {
+			p.count = 0
+		}
+	} else {
+		proxyIndex, _ = strconv.Atoi(s.Value)
+	}
+	p.proxy[proxyIndex].ServeHTTP(w, r)
+}
+
+// init httputil.ReverseProxy instance and add path rewrite and add session-sticky cookie to response
+func makeProxy(upstream string, pw model.PathRewrite, index int) *httputil.ReverseProxy {
+	p := &httputil.ReverseProxy{}
+	p.Director = func(req *http.Request) {
+		turl, _ := url.Parse(upstream)
+		req.URL.Host = turl.Host
+		req.URL.Scheme = turl.Scheme
+		if pw.Replace != "" {
+			req.URL.Path = strings.TrimPrefix(req.URL.Path, pw.Replace)
+			if req.URL.RawPath != "" {
+				req.URL.RawPath = strings.TrimPrefix(req.URL.RawPath, pw.Replace)
+			}
+		}
+		if pw.With != "" {
+			req.URL.Path = strings.TrimSuffix(pw.With, "/") + "/" + req.URL.Path
+			if req.URL.RawPath != "" {
+				req.URL.RawPath = strings.TrimSuffix(pw.With, "/") + "/" + req.URL.RawPath
+			}
+		}
+	}
+
+	p.ModifyResponse = func(resp *http.Response) error {
+		hasSticky := false
+		for _, cookie := range resp.Cookies() {
+			if cookie.Name == "s" {
+				hasSticky = true
+				break
+			}
+		}
+		if !hasSticky {
+			c := http.Cookie{
+				Name:  "s",
+				Value: strconv.Itoa(index),
+			}
+			resp.Header.Add("Set-Cookie", c.String())
+		}
+		return nil
+	}
+	return p
+}
+
+func NewProxyHandler(p *model.HttpPath) *ProxyHandler {
+	upstreamCount := len(p.Upstreams)
+	if upstreamCount == 0 {
+		panic("no upstream defined")
+	}
+	ph := &ProxyHandler{}
+	ph.proxy = make([]*httputil.ReverseProxy, upstreamCount)
+
+	for index, upstream := range p.Upstreams {
+		ph.proxy[index] = makeProxy(upstream, p.Rewrite, index)
+	}
+	return ph
+}
diff --git a/model/model.go b/model/model.go
index 1ad7b07..0a0785e 100644
--- a/model/model.go
+++ b/model/model.go
@@ -3,10 +3,16 @@ package model
 import "git.pyer.club/kingecg/gologger"
 
 type HttpPath struct {
-	Path         string `json:"path"`
-	Root         string `json:"root"`
-	Proxyto      string `json:"proxyto"`
-	ProxyHeaders string `json:"proxyheaders"`
+	Path      string      `json:"path"`
+	Root      string      `json:"root"`
+	Default   string      `json:"default"`
+	Upstreams []string    `json:"upstreams"`
+	Rewrite   PathRewrite `json:"pathrewrite"`
+}
+
+type PathRewrite struct {
+	Replace string `json:"replace"`
+	With    string `json:"with"`
 }
 
 type HttpServerConfig struct {