Zoraxy/mod/dynamicproxy/dynamicproxy.go

package dynamicproxy

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"imuslab.com/arozos/ReverseProxy/mod/dynamicproxy/dpcore"
	"imuslab.com/arozos/ReverseProxy/mod/reverseproxy"
	"imuslab.com/arozos/ReverseProxy/mod/tlscert"
)

/*
	Allow users to setup manual proxying for specific path

*/
type Router struct {
	ListenPort        int
	ProxyEndpoints    *sync.Map
	SubdomainEndpoint *sync.Map
	Running           bool
	Root              *ProxyEndpoint
	tlsCertManager    *tlscert.Manager
	mux               http.Handler
	TlsManager        *tlscert.Manager
	useTLS            bool
	server            *http.Server
	tlsListener       net.Listener
}

type RouterOption struct {
	Port int
}

type ProxyEndpoint struct {
	Root       string
	Domain     string
	RequireTLS bool
	Proxy      *dpcore.ReverseProxy `json:"-"`
}

type SubdomainEndpoint struct {
	MatchingDomain string
	Domain         string
	RequireTLS     bool
	Proxy          *reverseproxy.ReverseProxy `json:"-"`
}

type ProxyHandler struct {
	Parent *Router
}

func NewDynamicProxy(port int, useTls bool, tlsManager *tlscert.Manager) (*Router, error) {
	proxyMap := sync.Map{}
	domainMap := sync.Map{}
	thisRouter := Router{
		ListenPort:        port,
		ProxyEndpoints:    &proxyMap,
		SubdomainEndpoint: &domainMap,
		Running:           false,
		tlsCertManager:    tlsManager,
		useTLS:            useTls,
		server:            nil,
	}

	thisRouter.mux = &ProxyHandler{
		Parent: &thisRouter,
	}

	return &thisRouter, nil
}

//Update TLS setting in runtime. Will restart the proxy server
//if it is already running in the background
func (router *Router) UpdateTLSSetting(tlsEnabled bool) {
	router.useTLS = tlsEnabled
}

//Start the dynamic routing
func (router *Router) StartProxyService() error {
	//Create a new server object
	if router.server != nil {
		return errors.New("Reverse proxy server already running")
	}

	if router.Root == nil {
		return errors.New("Reverse proxy router root not set")
	}

	config := &tls.Config{
		GetCertificate: router.tlsCertManager.GetCert,
	}

	if router.useTLS {
		//Serve with TLS mode
		ln, err := tls.Listen("tcp", ":"+strconv.Itoa(router.ListenPort), config)
		if err != nil {
			log.Println(err)
			return err
		}
		router.tlsListener = ln
		router.server = &http.Server{Addr: ":" + strconv.Itoa(router.ListenPort), Handler: router.mux}
		router.Running = true

		if router.ListenPort == 443 {
			//Add a 80 to 443 redirector
			httpServer := &http.Server{
				Addr: ":80",
				Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					http.Redirect(w, r, "https://"+r.Host+r.RequestURI, http.StatusTemporaryRedirect)
				}),
				ReadTimeout:  3 * time.Second,
				WriteTimeout: 3 * time.Second,
				IdleTimeout:  120 * time.Second,
			}

			log.Println("Starting HTTP-to-HTTPS redirector (port 80)")
			go func() {
				//Start another router to check if the router.server is killed. If yes, kill this server as well
				go func() {
					for router.server != nil {
						time.Sleep(100 * time.Millisecond)
					}
					ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
					defer cancel()
					httpServer.Shutdown(ctx)
				}()
				if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
					log.Fatalf("Could not start server: %v\n", err)
				}
			}()
		}
		log.Println("Reverse proxy service started in the background (TLS mode)")
		go func() {
			if err := router.server.Serve(ln); err != nil && err != http.ErrServerClosed {
				log.Fatalf("Could not start server: %v\n", err)
			}
		}()
	} else {
		//Serve with non TLS mode
		router.tlsListener = nil
		router.server = &http.Server{Addr: ":" + strconv.Itoa(router.ListenPort), Handler: router.mux}
		router.Running = true
		log.Println("Reverse proxy service started in the background (Plain HTTP mode)")
		go func() {
			router.server.ListenAndServe()
			//log.Println("[DynamicProxy] " + err.Error())
		}()
	}

	return nil
}

func (router *Router) StopProxyService() error {
	if router.server == nil {
		return errors.New("Reverse proxy server already stopped")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := router.server.Shutdown(ctx)
	if err != nil {
		return err
	}

	if router.tlsListener != nil {
		router.tlsListener.Close()
	}

	//Discard the server object
	router.tlsListener = nil
	router.server = nil
	router.Running = false
	return nil
}

//Restart the current router if it is running.
//Startup the server if it is not running initially
func (router *Router) Restart() error {
	//Stop the router if it is already running
	if router.Running {
		err := router.StopProxyService()
		if err != nil {
			return err
		}
	}

	//Start the server
	err := router.StartProxyService()
	return err
}

/*
	Add an URL into a custom proxy services
*/
func (router *Router) AddVirtualDirectoryProxyService(rootname string, domain string, requireTLS bool) error {
	if domain[len(domain)-1:] == "/" {
		domain = domain[:len(domain)-1]
	}

	if rootname[len(rootname)-1:] == "/" {
		rootname = rootname[:len(rootname)-1]
	}

	webProxyEndpoint := domain
	if requireTLS {
		webProxyEndpoint = "https://" + webProxyEndpoint
	} else {
		webProxyEndpoint = "http://" + webProxyEndpoint
	}
	//Create a new proxy agent for this root
	path, err := url.Parse(webProxyEndpoint)
	if err != nil {
		return err
	}

	proxy := dpcore.NewDynamicProxyCore(path, rootname)

	endpointObject := ProxyEndpoint{
		Root:       rootname,
		Domain:     domain,
		RequireTLS: requireTLS,
		Proxy:      proxy,
	}

	router.ProxyEndpoints.Store(rootname, &endpointObject)

	log.Println("Adding Proxy Rule: ", rootname+" to "+domain)
	return nil
}

/*
	Remove routing from RP

*/
func (router *Router) RemoveProxy(ptype string, key string) error {
	fmt.Println(ptype, key)
	if ptype == "vdir" {
		router.ProxyEndpoints.Delete(key)
		return nil
	} else if ptype == "subd" {
		router.SubdomainEndpoint.Delete(key)
		return nil
	}
	return errors.New("invalid ptype")
}

/*
	Add an default router for the proxy server
*/
func (router *Router) SetRootProxy(proxyLocation string, requireTLS bool) error {
	if proxyLocation[len(proxyLocation)-1:] == "/" {
		proxyLocation = proxyLocation[:len(proxyLocation)-1]
	}

	webProxyEndpoint := proxyLocation
	if requireTLS {
		webProxyEndpoint = "https://" + webProxyEndpoint
	} else {
		webProxyEndpoint = "http://" + webProxyEndpoint
	}
	//Create a new proxy agent for this root
	path, err := url.Parse(webProxyEndpoint)
	if err != nil {
		return err
	}

	proxy := dpcore.NewDynamicProxyCore(path, "")

	rootEndpoint := ProxyEndpoint{
		Root:       "/",
		Domain:     proxyLocation,
		RequireTLS: requireTLS,
		Proxy:      proxy,
	}

	router.Root = &rootEndpoint
	return nil
}

//Do all the main routing in here
func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if strings.Contains(r.Host, ".") {
		//This might be a subdomain. See if there are any subdomain proxy router for this
		//Remove the port if any
		domainOnly := r.Host
		if strings.Contains(r.Host, ":") {
			hostPath := strings.Split(r.Host, ":")
			domainOnly = hostPath[0]
		}
		sep := h.Parent.getSubdomainProxyEndpointFromHostname(domainOnly)
		if sep != nil {
			h.subdomainRequest(w, r, sep)
			return
		}
	}

	//Clean up the request URI
	proxyingPath := strings.TrimSpace(r.RequestURI)
	/*
		requestingRsc := filepath.Base(r.RequestURI)
		if !strings.Contains(requestingRsc, ".") {
			//Not a file. Add the tail slash if not exists
			if proxyingPath[len(proxyingPath)-1:] != "/" {
				proxyingPath = proxyingPath + "/"
				http.Redirect(w, r, proxyingPath, http.StatusTemporaryRedirect)
				return
			}
		}
	*/

	targetProxyEndpoint := h.Parent.getTargetProxyEndpointFromRequestURI(proxyingPath)
	if targetProxyEndpoint != nil {
		h.proxyRequest(w, r, targetProxyEndpoint)
	} else {
		h.proxyRequest(w, r, h.Parent.Root)
	}
}