認証付きプロキシを超えるためのプロキシを書いた

職場が認証付きプロキシを越えないと外に出ていけない環境なのだけど、ツールによって認証付きプロキシのための設定が面倒だったり、そもそも無理だったりして何かとうっとうしいので、リクエストに認証情報を付加するプロキシを書いた。
こいつをローカルで動かして、各ツールのプロキシには localhost:1080 を指定する。

squidとかでできるでしょという話もあるんだけど、squidの設定ファイル書くよりコード書くほうが楽しいし。

proxy.go

package main

import (
    "encoding/base64"
    "fmt"
    "io"
    "log"
    "net"
    "net/http"
    "net/url"
)

var (
    port               = 1080
    proxyHost          = "proxy.example.com:3128"
    proxyUser          = "iwata"
    proxyPass          = "mypassword"
    proxyAuthorization = "Basic " + base64.StdEncoding.EncodeToString([]byte(proxyUser+":"+proxyPass))
)

func handleHttps(w http.ResponseWriter, r *http.Request) {
    hj, _ := w.(http.Hijacker)
    if proxyConn, err := net.Dial("tcp", proxyHost); err != nil {
        log.Fatal(err)
    } else if clientConn, _, err := hj.Hijack(); err != nil {
        proxyConn.Close()
        log.Fatal(err)
    } else {
        r.Header.Set("Proxy-Authorization", proxyAuthorization)
        r.Write(proxyConn)
        go func() {
            io.Copy(clientConn, proxyConn)
            proxyConn.Close()
        }()
        go func() {
            io.Copy(proxyConn, clientConn)
            clientConn.Close()
        }()
    }
}

func handleHttp(w http.ResponseWriter, r *http.Request) {
    hj, _ := w.(http.Hijacker)
    client := &http.Client{}
    r.RequestURI = ""
    if resp, err := client.Do(r); err != nil {
        log.Fatal(err)
    } else if conn, _, err := hj.Hijack(); err != nil {
        log.Fatal(err)
    } else {
        defer conn.Close()
        resp.Write(conn)
    }
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
    log.Printf("[%s] %s", r.Method, r.URL)
    if r.Method == "CONNECT" {
        handleHttps(w, r)
    } else {
        handleHttp(w, r)
    }
}

func main() {
    proxyUrlString := fmt.Sprintf("http://%s:%s@%s", url.QueryEscape(proxyUser), url.QueryEscape(proxyPass), proxyHost)
    proxyUrl, err := url.Parse(proxyUrlString)
    if err != nil {
        log.Fatal(err)
    }
    http.DefaultTransport = &http.Transport{Proxy: http.ProxyURL(proxyUrl)}

    handler := http.HandlerFunc(handleRequest)
    listen := fmt.Sprintf("localhost:%d", port)
    log.Printf("Start serving on %s", listen)
    log.Fatal(http.ListenAndServe(listen, handler))
}

Trivial HTTP Proxy in Go という記事があるけど、今のバージョン(go 1.3)ではこのコードはちゃんと動かなかった。
ResponseWriterに直接Writeすると、ステータスコードやヘッダも全部ボディに書き込まれてしまうようで、Hijack()で生のコネクションを取り出してから書き込む必要があるっぽい。