init

2022-05-17 01:56:44 +03:00
commit 3a9b8c8468
117 changed files with 7175 additions and 0 deletions
@@ -0,0 +1,121 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+	"time"
+)
+
+var (
+	listen      string
+	gracePeriod time.Duration
+	target      string
+	queue       = make(chan struct{}, 1)
+)
+
+func init() {
+	flag.StringVar(&listen, "listen", ":4444", "host and port to listen to")
+	flag.DurationVar(&gracePeriod, "grace-period", 300*time.Second, "graceful shutdown period")
+	flag.StringVar(&target, "target", "http://localhost:4723", "target proxy url")
+	flag.Parse()
+}
+
+func signalContext() (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		defer cancel()
+		stop := make(chan os.Signal)
+		signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
+		<-stop
+	}()
+	return ctx, cancel
+}
+
+func main() {
+	ctx, cancel := signalContext()
+	defer cancel()
+	err := mainCtx(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func mux(ctx context.Context, u *url.URL) http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/wd/hub/session", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-r.Context().Done():
+			return
+		case <-ctx.Done():
+			http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
+			return
+		case queue <- struct{}{}:
+		}
+		(&httputil.ReverseProxy{
+			Director: func(r *http.Request) {
+				r.URL.Scheme, r.URL.Host = u.Scheme, u.Host
+			},
+			ModifyResponse: func(resp *http.Response) error {
+				if resp.StatusCode != http.StatusOK {
+					<-queue
+				}
+				return nil
+			},
+			ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
+				<-queue
+				log.Printf("proxy to %s: %v", u.String(), err)
+				http.Error(w, http.StatusText(http.StatusBadGateway), http.StatusBadGateway)
+
+			},
+		}).ServeHTTP(w, r)
+	}))
+	mux.HandleFunc("/wd/hub/session/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		(&httputil.ReverseProxy{
+			Director: func(r *http.Request) {
+				r.URL.Scheme, r.URL.Host = u.Scheme, u.Host
+			},
+		}).ServeHTTP(w, r)
+		fragments := strings.Split(r.URL.Path, "/")
+		if r.Method == http.MethodDelete && len(fragments) == 5 {
+			<-queue
+		}
+	}))
+	return mux
+}
+
+func mainCtx(ctx context.Context) error {
+	target, err := url.Parse(target)
+	if err != nil {
+		return fmt.Errorf("parse target url: %v", err)
+	}
+	server := &http.Server{
+		Addr:    listen,
+		Handler: mux(ctx, target),
+	}
+	e := make(chan error)
+	go func() {
+		e <- server.ListenAndServe()
+	}()
+	select {
+	case err := <-e:
+		return err
+	case <-ctx.Done():
+	}
+	log.Printf("starting graceful shutdown in %v]", gracePeriod)
+	shCtx, cancel := context.WithTimeout(context.Background(), gracePeriod)
+	defer cancel()
+	if err := server.Shutdown(shCtx); err != nil {
+		return fmt.Errorf("graceful shutdown: %v", err)
+	}
+	log.Printf("stopped")
+	return nil
+}