122 lines
2.8 KiB
Go
122 lines
2.8 KiB
Go
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
|
|
}
|