Совместное использование контекста между HandleFunc и modifyResponse в Go reverse proxy

#go #proxy #reverse-proxy

#Вперед #прокси #обратный прокси

Вопрос:

Я пытаюсь изучить Go и подумал, что неплохим небольшим проектом будет A / B тестовый прокси для размещения перед веб-сервером. Я и не подозревал, что Go, по сути, предлагает обратный прокси из коробки, поэтому настройка была простой. Я дошел до того, что проксирую трафик, но вот в чем дело, у меня возникают проблемы с реализацией фактической функциональности, потому что везде, где у меня есть доступ к ответу, у меня нет доступа к назначенным вариантам A / B теста:

  • В handleFunc я назначаю варианты каждого теста запросу, чтобы вышестоящий сервер также мог знать об этом и использовать if для реализаций в своем бэкэнде.
  • Я устанавливаю cookie со всеми тестами и вариациями, как для запроса, который передается через прокси-сервер в вышестоящий канал, так и для ответа, который возвращается клиенту.
  • Тесты, состоящие из пары поиск / замена, будут вносить изменения в тело ответа после того, как ответ вернется с вышестоящего сервера.
  • Я пытаюсь использовать modifyResponse функцию httputil.ReverseProxy для изменения ответа.

Проблема в том, что я не могу понять, как совместно использовать назначенные варианты между handleFunc и modifyResponse , не меняя вышестоящий сервер. Я хотел бы иметь возможность поделиться этим контекстом (в основном map[string]string каким-то образом.

Пример кода:

Вот дистиллированная версия моего кода, где мой вопрос в основном заключается в том, как можно modifyRequest узнать о случайных назначениях, которые произошли в handleFunc ?

 package main

import (
    config2 "ab-proxy/config"
    "bytes"
    "fmt"
    "io/ioutil"
    "net/http"
    "net/http/httputil"
    "net/url"
    "strconv"
    "strings"
)

var config config2.ProxyConfig
var reverseProxy *httputil.ReverseProxy
var tests config2.Tests

func overwriteCookie(req *http.Request, cookie *http.Cookie) {
    // omitted for brevity, will replace a cookie header, instead of adding a second value
}

func parseRequestCookiesToAssignedTests(req *http.Request) map[string]string {
    // omitted for brevity, builds a map where the key is the identifier of the test, the value the assigned variant
}

func renderCookieForAssignedTests(assignedTests map[string]string) string {
    // omitted for brevity, builds a cookie string
}

func main () {
    var err error

    if  config, err = config2.LoadConfig(); err != nil {
        fmt.Println(err)

        return
    }

    if tests, err = config2.LoadTests(); err != nil {
        fmt.Println(err)

        return
    }

    upstreamUrl, _ := url.Parse("0.0.0.0:80")

    reverseProxy = httputil.NewSingleHostReverseProxy(upstreamUrl)
    reverseProxy.ModifyResponse = modifyResponse

    http.HandleFunc("/", handleRequest)

    if err := http.ListenAndServe("0.0.0.0:80", nil); err != nil {
        fmt.Println("Could not start proxy")
    }
}

func handleRequest(res http.ResponseWriter, req *http.Request) {
    assigned := parseRequestCookiesToAssignedTests(req)

    newCookies := make(map[string]string)

    for _, test := range tests.Entries {
        val, ok := assigned[test.Identifier]

        if ok {
            newCookies[test.Identifier] = val
        } else {
            newCookies[test.Identifier] = "not-assigned-yet" // this will be replaced by random variation assignment
        }
    }

    testCookie := http.Cookie{Name: config.Cookie.Name, Value: renderCookieForAssignedTests(newCookies)}

    // Add cookie to request to be sent to upstream
    overwriteCookie(req, amp;testCookie)

    // Add cookie to response to be returned to client
    http.SetCookie(res, amp;testCookie)

    reverseProxy.ServeHTTP(res, req)
}

func modifyResponse (response *http.Response) error {
    body, err := ioutil.ReadAll(response.Body)

    if err != nil {
        return  err
    }

    err = response.Body.Close()

    if err != nil {
        return err
    }

    response.Body = ioutil.NopCloser(bytes.NewReader(body))
    response.ContentLength = int64(len(body))
    response.Header.Set("Content-Length", strconv.Itoa(len(body)))

    return nil
}
  

Ответ №1:

Используйте стандартный context.Context . Это доступно в вашем обработчике через *http.Request . И запрос также доступен через *http.Response аргумент to modifyResponse .

В вашем обработчике:

 ctx := req.Context()
// Set values, deadlines, etc.
req = req.WithContext(ctx)
reverseProxy.ServeHTTP(res, req)
  

Затем в modifyResponse :

 ctx := response.Request.Context()
// fetch values, check for cancellation, etc
  

Комментарии:

1. Ах, потрясающе! Даже навороченная среда разработки не смогла бы помочь мне выяснить, что у .Request есть response свойство и что на самом деле существует универсальный context.Context класс, который делает именно то, что я ищу. Это очень помогает, спасибо!