// Will receive an input stream which would convert the response to utf-8
// The given function must close the reader r, in order to close the response body.
func HandleStringReader(f func(r io.Reader, ctx *goproxy.ProxyCtx) io.Reader) goproxy.RespHandler {
return goproxy.FuncRespHandler(func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
if ctx.Error != nil {
return nil
}
charsetName := ctx.Charset()
if charsetName == "" {
charsetName = "utf-8"
}
if strings.ToLower(charsetName) != "utf-8" {
r, err := charset.NewReader(charsetName, resp.Body)
if err != nil {
ctx.Warnf("Cannot convert from %v to utf-8: %v", charsetName, err)
return resp
}
tr, err := charset.TranslatorTo(charsetName)
if err != nil {
ctx.Warnf("Can't translate to %v from utf-8: %v", charsetName, err)
return resp
}
if err != nil {
ctx.Warnf("Cannot translate to %v: %v", charsetName, err)
return resp
}
newr := charset.NewTranslatingReader(f(r, ctx), tr)
resp.Body = &readFirstCloseBoth{ioutil.NopCloser(newr), resp.Body}
} else {
//no translation is needed, already at utf-8
resp.Body = &readFirstCloseBoth{ioutil.NopCloser(f(resp.Body, ctx)), resp.Body}
}
return resp
})
}
func translate(tr charset.Translator, in string) (string, error) {
var buf bytes.Buffer
r := charset.NewTranslatingReader(strings.NewReader(in), tr)
_, err := io.Copy(&buf, r)
if err != nil {
return "", err
}
return string(buf.Bytes()), nil
}
func testTranslatingReader(t *testing.T, tr charset.Translator, inReader, outReader func(io.Reader) io.Reader) {
data := make([]byte, 128)
for i := range data {
data[i] = byte(i)
}
inr := inReader(bytes.NewBuffer(data))
r := charset.NewTranslatingReader(inr, tr)
outr := outReader(r)
r = outr
var outbuf bytes.Buffer
_, err := io.Copy(&outbuf, r)
if err != nil {
t.Fatalf("translator %T, readers %T, %T, copy failed: %v", tr, inr, outr, err)
}
err = checkTranslation(data, outbuf.Bytes())
if err != nil {
t.Fatalf("translator %T, readers %T, %T, %v\n", err)
}
}