151 lines
3.8 KiB
Go
151 lines
3.8 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"io/ioutil"
|
|
"log"
|
|
"net/http"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"os"
|
|
"os/signal"
|
|
"path"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
listen = ":4444"
|
|
target = "http://localhost:5555"
|
|
waitTimeout = 30 * time.Second
|
|
gracePeriod = 30 * time.Second
|
|
browserName = "safari"
|
|
browserVersion = "15.0"
|
|
)
|
|
|
|
func wait(ctx context.Context, target string) (*url.URL, error) {
|
|
for {
|
|
r, err := http.NewRequest(http.MethodHead, target, http.NoBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("new %s request to %s: %v", http.MethodHead, target, err)
|
|
}
|
|
resp, err := http.DefaultClient.Do(r.WithContext(ctx))
|
|
if resp != nil {
|
|
resp.Body.Close()
|
|
}
|
|
if err != nil {
|
|
if err, ok := err.(*url.Error); ok {
|
|
switch err.Err {
|
|
case context.Canceled, context.DeadlineExceeded:
|
|
return nil, err
|
|
default:
|
|
<-time.After(100 * time.Millisecond)
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
return r.URL, nil
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
e := make(chan error)
|
|
go func() {
|
|
stop := make(chan os.Signal)
|
|
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
|
|
select {
|
|
case err := <-e:
|
|
log.Fatalf("server: %v", err)
|
|
case <-stop:
|
|
cancel()
|
|
}
|
|
}()
|
|
waitCtx, waitCancel := context.WithTimeout(ctx, waitTimeout)
|
|
defer waitCancel()
|
|
u, err := wait(waitCtx, target)
|
|
if err != nil {
|
|
log.Fatal(fmt.Errorf("wait target: %v", err))
|
|
}
|
|
server := &http.Server{
|
|
Addr: listen,
|
|
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var value map[string]interface{}
|
|
err := json.NewDecoder(r.Body).Decode(&value)
|
|
if err != nil && err != io.EOF {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if err == nil {
|
|
if _, ok := value["desiredCapabilities"]; ok {
|
|
delete(value, "desiredCapabilities")
|
|
}
|
|
if o, ok := value["capabilities"]; ok {
|
|
if w3cCapabilities, ok := o.(map[string]interface{}); ok {
|
|
for _, match := range []string{"alwaysMatch", "firstMatch"} {
|
|
delete(w3cCapabilities, match)
|
|
}
|
|
}
|
|
}
|
|
body, err := json.Marshal(value)
|
|
if err != nil {
|
|
log.Printf("[ERROR] marshalling capabilities: %v", err)
|
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
r.Body = ioutil.NopCloser(bytes.NewReader(body))
|
|
r.ContentLength = int64(len(body))
|
|
}
|
|
(&httputil.ReverseProxy{
|
|
Director: func(r *http.Request) {
|
|
r.URL.Scheme, r.URL.Host, r.URL.Path = u.Scheme, u.Host, path.Join(u.Path, r.URL.Path)
|
|
},
|
|
ModifyResponse: func(resp *http.Response) error {
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil
|
|
}
|
|
var values map[string]interface{}
|
|
defer resp.Body.Close()
|
|
err := json.NewDecoder(resp.Body).Decode(&values)
|
|
if err != nil {
|
|
return fmt.Errorf("decode json response: %v", err)
|
|
}
|
|
if o, ok := values["value"]; ok {
|
|
if value, ok := o.(map[string]interface{}); ok {
|
|
if o, ok := value["capabilities"]; ok {
|
|
if capabilities, ok := o.(map[string]interface{}); ok {
|
|
capabilities["browserName"] = browserName
|
|
capabilities["browserVersion"] = browserVersion
|
|
delete(capabilities, "platformName")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
buf, err := json.Marshal(&values)
|
|
if err != nil {
|
|
return fmt.Errorf("encode json response: %v", err)
|
|
}
|
|
resp.Header.Del("Server")
|
|
resp.Header.Del("Content-Length")
|
|
resp.ContentLength = int64(len(buf))
|
|
resp.Body = ioutil.NopCloser(bytes.NewReader(buf))
|
|
return nil
|
|
},
|
|
}).ServeHTTP(w, r)
|
|
}),
|
|
}
|
|
go func() {
|
|
e <- server.ListenAndServe()
|
|
}()
|
|
<-ctx.Done()
|
|
shCtx, shCancel := context.WithTimeout(context.Background(), gracePeriod)
|
|
defer shCancel()
|
|
if err := server.Shutdown(shCtx); err != nil {
|
|
log.Fatalf("graceful shutdown: %v]", err)
|
|
}
|
|
}
|