[Go] rubygemsをキャッシュするプロキシサーバのサンプル

とりあえず動いたというレベルのサンプルですが、rubygemsをキャッシュするプロキシサーバのサンプルを2種類書いてみました。
実行例はgithubのプロジェクトページを参照してください。

hnakamur/goproxy-rubygems-cache

elazarl/goproxyを使った実装

package main

import (
    "flag"
    "github.com/elazarl/goproxy"
    "io/ioutil"
    "log"
    "net/http"
    "os"
    "path"
)

func main() {
    verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
    docRoot := flag.String("root", ".", "document root directory")
    address := flag.String("http", ":8080", `HTTP service address (e.g., ":8080")`)
    flag.Parse()
    proxy := goproxy.NewProxyHttpServer()
    proxy.Verbose = *verbose

    proxy.OnRequest(reqMethodIs("GET", "HEAD")).DoFunc(
        func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
            filename := path.Join(*docRoot, ctx.Req.URL.Path)
            if !exists(filename) {
                return req, nil
            }

            bytes, err := ioutil.ReadFile(filename)
            if err != nil {
                ctx.Warnf("%s", err)
                return req, nil
            }
            resp := goproxy.NewResponse(req, "application/octet-stream",
                http.StatusOK, string(bytes))
            ctx.Logf("return response from local %s", filename)
            return req, resp
        })

    proxy.OnResponse(respReqMethodIs("GET", "HEAD")).Do(
        goproxy.HandleBytes(
            func(b []byte, ctx *goproxy.ProxyCtx) []byte {
                if ctx.Req.Method != "GET" || hasRespHeader(ctx.Resp, "Location") {
                    return b
                }

                filename := path.Join(*docRoot, ctx.Req.URL.Path)
                if exists(filename) {
                    return b
                }

                dir := path.Dir(filename)
                err := os.MkdirAll(dir, 0755)
                if err != nil {
                    ctx.Warnf("cannot create directory: %s", dir)
                }

                err = ioutil.WriteFile(filename, b, 0644)
                if err != nil {
                    ctx.Warnf("cannot write file: %s", filename)
                }

                ctx.Logf("save cache to %s", filename)

                return b
            }))
    log.Fatal(http.ListenAndServe(*address, proxy))
}

func reqMethodIs(methods ...string) goproxy.ReqConditionFunc {
    return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
        for _, method := range methods {
            if req.Method == method {
                return true
            }
        }
        return false
    }
}

func respReqMethodIs(methods ...string) goproxy.RespConditionFunc {
    return func(resp *http.Response, ctx *goproxy.ProxyCtx) bool {
        for _, method := range methods {
            if resp.Request.Method == method {
                return true
            }
        }
        return false
    }
}

func hasRespHeader(resp *http.Response, header string) bool {
    _, ok := resp.Header[header]
    return ok
}

func exists(filename string) bool {
    _, err := os.Stat(filename)
    return !os.IsNotExist(err)
}

hnakamur/rubygems-proxy

Trivial HTTP Proxy in Go — IOTTMCOを参考にGo標準ライブラリを使った実装

package main

import (
    "bufio"
    "flag"
    "io"
    "log"
    "net/http"
    "os"
    "path"
    "strings"
)

var addr = flag.String("addr", ":8080", "http service address")
var root = flag.String("root", ".", "document root directory")

func main() {
    flag.Parse()
    http.Handle("/", http.HandlerFunc(handler))
    err := http.ListenAndServe(*addr, nil)
    if err != nil {
        log.Fatal("ListenAndServe:", err)
    }
}

func handler(w http.ResponseWriter, req *http.Request) {
    log.Printf("method:%s\turl:%s", req.Method, req.URL)
    filename := path.Join(*root, req.URL.Path)
    if exists(filename) {
        http.ServeFile(w, req, filename)
        log.Printf("served from local file: %s", filename)
    } else {
        proxy(w, req)
    }
}

func proxy(w http.ResponseWriter, req *http.Request) {
    client := &http.Client{}

    // Tweak the request as appropriate:
    //  RequestURI may not be sent to client
    //  URL.Scheme must be lower-case
    req.RequestURI = ""
    req.URL.Scheme = strings.ToLower(req.URL.Scheme)

    // And proxy
    resp, err := client.Do(req)
    if err != nil {
        log.Fatal(err)
    }

    if req.Method != "GET" || hasLocationHeader(resp) {
        log.Printf("forward proxy to remote: %s", req.URL)
        resp.Write(w)
        return
    }

    filename := path.Join(*root, req.URL.Path)
    dir := path.Dir(filename)
    if !exists(dir) {
        os.MkdirAll(dir, 0755)
    }

    file, err := os.Create(filename)
    if err != nil {
        log.Fatal(err)
    }

    fw := bufio.NewWriter(file)
    _, err = io.Copy(fw, resp.Body)
    if err != nil {
        log.Fatal(err)
    }
    err = resp.Body.Close()
    if err != nil {
        log.Fatal(err)
    }
    err = fw.Flush()
    if err != nil {
        log.Fatal(err)
    }
    err = file.Close()
    if err != nil {
        log.Fatal(err)
    }
    log.Printf("saved to local file: %s", filename)

    http.ServeFile(w, req, filename)
}

func hasLocationHeader(resp *http.Response) bool {
    _, ok := resp.Header["Location"]
    return ok
}

func exists(filename string) bool {
    _, err := os.Stat(filename)
    return !os.IsNotExist(err)
}

個人的にはこちらのほうがわかりやすい気がします。