用 Golang 建立一個簡單的 Reverse Proxy

前言

雷雷寫上一篇文章的時候發現部落格的音樂 API 失效了，發現原本的作者有給新的 API ，但有限制一些 Header 才能用

於是雷雷就想寫一個 Reverse Proxy 幫忙轉發請求 ay

因為要放在 Docker 裡面，想了想還是 binary 好處理，於是把目標放在 Golang 上

需求

• Golang

實作

核心原始碼

func main() {

    url, err := url.Parse("https://" + *remoteHost)
    if err != nil {
        panic(err)
    }

    director := func(req *http.Request) {
        req.URL.Scheme = url.Scheme
        req.URL.Host = url.Host
        req.Host = url.Host
        req.Header.Set("Referer", "https://"+*referrer+"/")
        req.Header.Set("Origin", "https://"+*referrer)
        req.Header.Set("authority", *remoteHost)
    }

    reverseProxy := httputil.NewSingleHostReverseProxy(url)
    reverseProxy.Director = director
    handler := handler{proxy: reverseProxy}
    http.Handle("/", handler)

    http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)
}

type handler struct {
    proxy *httputil.ReverseProxy
}

func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    h.proxy.ServeHTTP(w, r)
}

這些就是最核心的 Code 了

可以看到我們先用 httputil.NewSingleHostReverseProxy 建一個 Reverse Proxy 起來

用 director 干涉 Income Request 後轉發到目標 (url) 上。

干涉回應

除此之外，如果我們想要修改 Response Body 的某一部分的話，我們可以用 ModifyResponse 來實作:

    reverseProxy := httputil.NewSingleHostReverseProxy(url)
    reverseProxy.Director = director
    reverseProxy.ModifyResponse = rewriteBody
    // Else Code...
}

func rewriteBody(resp *http.Response) (err error) {
    b, err := ioutil.ReadAll(resp.Body)
    if err != nil {
        return err
    }

    err = resp.Body.Close()
    if err != nil {
        return err
    }

    b = bytes.Replace(b, []byte("before"), []byte("after", -1)
    body := ioutil.NopCloser(bytes.NewReader(b))
    resp.Body = body
    resp.ContentLength = int64(len(b))
    resp.Header.Set("Content-Length", strconv.Itoa(len(b)))
    return nil
}

但這樣還不夠: 如果伺服器端開啟了壓縮(例如 gzip)，bytes 就會讀不到
因此我們要根據情況解壓縮:

    var reader io.ReadCloser
    switch resp.Header.Get("Content-Encoding") {
    case "gzip":
        reader, err = gzip.NewReader(resp.Body)
        defer reader.Close()
    default:
        reader = resp.Body
    }
    delete(resp.Header, "Content-Encoding")

    b, err := ioutil.ReadAll(reader)

因此我們的 Code 最後可以像這樣:

package main

import (
    "bytes"
    "compress/gzip"
    "flag"
    "fmt"
    "io"
    "io/ioutil"
    "net/http"
    "net/http/httputil"
    "net/url"
    "strconv"
)

var remoteHost *string
var referrer *string
var myApiHost *string
var myApiHostSchema *string

func main() {
    port := flag.Int("p", 80, "HTTP API Port")
    remoteHost = flag.String("r", "target.host", "Remote Host")
    referrer = flag.String("f", "target.host", "Referrer")
    myApiHost = flag.String("m", "localhost", "My API Host")
    myApiHostSchema = flag.String("s", "http", "My API Host Schema")
    flag.Parse()

    url, err := url.Parse("https://" + *remoteHost)
    if err != nil {
        panic(err)
    }

    director := func(req *http.Request) {
        req.URL.Scheme = url.Scheme
        req.URL.Host = url.Host
        req.Host = url.Host
        req.Header.Set("Referer", "https://"+*referrer+"/")
        req.Header.Set("Origin", "https://"+*referrer)
        req.Header.Set("authority", *remoteHost)
    }

    reverseProxy := httputil.NewSingleHostReverseProxy(url)
    reverseProxy.Director = director
    reverseProxy.ModifyResponse = rewriteBody
    handler := handler{proxy: reverseProxy}
    http.Handle("/", handler)

    http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)
}

type handler struct {
    proxy *httputil.ReverseProxy
}

func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    h.proxy.ServeHTTP(w, r)
}

func rewriteBody(resp *http.Response) (err error) {
    var reader io.ReadCloser
    switch resp.Header.Get("Content-Encoding") {
    case "gzip":
        reader, err = gzip.NewReader(resp.Body)
        defer reader.Close()
    default:
        reader = resp.Body
    }
    delete(resp.Header, "Content-Encoding")

    b, err := ioutil.ReadAll(reader)

    if err != nil {
        return err
    }

    err = resp.Body.Close()
    if err != nil {
        return err
    }
    b = bytes.Replace(b, []byte("https:\\/\\/"+*remoteHost), []byte(*myApiHostSchema+":\\/\\/"+*myApiHost), -1)
    body := ioutil.NopCloser(bytes.NewReader(b))
    resp.Body = body
    resp.ContentLength = int64(len(b))
    resp.Header.Set("Content-Length", strconv.Itoa(len(b)))
    return nil
}

簡單又實用的一小段 Golang Code，耶比!