117 lines
2.6 KiB
Go
117 lines
2.6 KiB
Go
package broker
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"local/storage"
|
|
"local/truckstop/config"
|
|
"local/truckstop/logtr"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// once per minute
|
|
var authlimiter = rate.NewLimiter(rate.Limit(1.0/60.0), 1)
|
|
|
|
// thrice per minute
|
|
var limiter = rate.NewLimiter(rate.Limit(1.0/20.0), 1)
|
|
|
|
type Broker interface {
|
|
Search([]config.State) ([]Job, error)
|
|
}
|
|
|
|
func do(r *http.Request) (*http.Response, error) {
|
|
return _do(config.Get().DB(), r)
|
|
}
|
|
|
|
func _do(db storage.DB, r *http.Request) (*http.Response, error) {
|
|
limiter.Wait(context.Background())
|
|
if strings.Contains(strings.ToLower(r.URL.Path), "login") {
|
|
authlimiter.Wait(context.Background())
|
|
}
|
|
client := &http.Client{
|
|
Timeout: time.Hour,
|
|
}
|
|
|
|
cookieJarKey := "cookies_" + r.URL.Host
|
|
cookies, err := getCookies(db, cookieJarKey)
|
|
if err != nil {
|
|
logtr.Errorf("failed to get cookies: %v", err)
|
|
} else {
|
|
logtr.Verbosef("got cookies for %s: %+v", cookieJarKey, cookies)
|
|
cookieV := strings.Join(r.Header["Cookie"], "; ")
|
|
for _, cookie := range cookies {
|
|
if len(cookieV) > 0 {
|
|
cookieV += "; "
|
|
}
|
|
cookieV += cookie
|
|
}
|
|
r.Header.Set("Cookie", cookieV)
|
|
}
|
|
logtr.Verbosef("_do: %+v", r)
|
|
resp, err := client.Do(r)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err := setCookies(db, cookieJarKey, resp); err != nil {
|
|
logtr.Errorf("failed to set cookies: %v", err)
|
|
}
|
|
return resp, err
|
|
}
|
|
|
|
func getCookies(db storage.DB, host string) ([]string, error) {
|
|
b, err := db.Get(host)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var result []string
|
|
err = json.Unmarshal(b, &result)
|
|
return result, err
|
|
}
|
|
|
|
func setCookies(db storage.DB, host string, resp *http.Response) error {
|
|
nameValues := [][2]string{}
|
|
old, _ := getCookies(db, host)
|
|
for _, value := range old {
|
|
if len(value) == 0 {
|
|
continue
|
|
}
|
|
name := strings.Split(value, "=")[0]
|
|
value := strings.Split(value, ";")[0]
|
|
nameValues = append(nameValues, [2]string{name, value})
|
|
}
|
|
for _, cookie := range resp.Header["Set-Cookie"] {
|
|
if len(cookie) == 0 {
|
|
continue
|
|
}
|
|
name := strings.Split(cookie, "=")[0]
|
|
value := strings.Split(cookie, ";")[0]
|
|
found := false
|
|
for i := range nameValues {
|
|
if nameValues[i][0] == name {
|
|
found = true
|
|
nameValues[i][1] = value
|
|
}
|
|
}
|
|
if !found {
|
|
nameValues = append(nameValues, [2]string{name, value})
|
|
}
|
|
}
|
|
result := []string{}
|
|
for i := range nameValues {
|
|
if len(nameValues[i][1]) == 0 {
|
|
continue
|
|
}
|
|
result = append(result, nameValues[i][1])
|
|
}
|
|
logtr.Verbosef("setting cookies for %s: %+v", host, result)
|
|
b, err := json.Marshal(result)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return db.Set(host, b)
|
|
}
|