vhost: a bug fix of reading request

This commit is contained in:
fatedier 2017-07-01 12:13:44 +08:00
parent 4cc1663a5f
commit 0752508469

View File

@ -76,8 +76,12 @@ func hostNameRewrite(request io.Reader, rewriteHost string, remoteIP string) (_
buf := pool.GetBuf(1024) buf := pool.GetBuf(1024)
defer pool.PutBuf(buf) defer pool.PutBuf(buf)
request.Read(buf) var n int
retBuffer, err := parseRequest(buf, rewriteHost, remoteIP) n, err = request.Read(buf)
if err != nil {
return
}
retBuffer, err := parseRequest(buf[:n], rewriteHost, remoteIP)
return retBuffer, err return retBuffer, err
} }
@ -113,8 +117,8 @@ func parseRequest(org []byte, rewriteHost string, remoteIP string) (ret []byte,
} }
buf := new(bytes.Buffer) buf := new(bytes.Buffer)
buf.Write(b) buf.Write(b)
buf.WriteString(fmt.Sprintf("X-Forwarded-For: %s\n", remoteIP)) buf.WriteString(fmt.Sprintf("X-Forwarded-For: %s\r\n", remoteIP))
buf.WriteString(fmt.Sprintf("X-Real-IP: %s\n", remoteIP)) buf.WriteString(fmt.Sprintf("X-Real-IP: %s\r\n", remoteIP))
if len(changedBuf) == 0 { if len(changedBuf) == 0 {
tp.WriteTo(buf) tp.WriteTo(buf)
} else { } else {
@ -138,8 +142,8 @@ func parseRequest(org []byte, rewriteHost string, remoteIP string) (ret []byte,
firstLine := req.Method + " " + req.URL.String() + " " + req.Proto firstLine := req.Method + " " + req.URL.String() + " " + req.Proto
buf := new(bytes.Buffer) buf := new(bytes.Buffer)
buf.WriteString(firstLine) buf.WriteString(firstLine)
buf.WriteString(fmt.Sprintf("X-Forwarded-For: %s\n", remoteIP)) buf.WriteString(fmt.Sprintf("X-Forwarded-For: %s\r\n", remoteIP))
buf.WriteString(fmt.Sprintf("X-Real-IP: %s\n", remoteIP)) buf.WriteString(fmt.Sprintf("X-Real-IP: %s\r\n", remoteIP))
tp.WriteTo(buf) tp.WriteTo(buf)
return buf.Bytes(), err return buf.Bytes(), err
} }
@ -175,9 +179,9 @@ func changeHostName(buff *bytes.Buffer, rewriteHost string) (_ []byte, err error
var hostHeader string var hostHeader string
portPos := bytes.IndexByte(kv[j+1:], ':') portPos := bytes.IndexByte(kv[j+1:], ':')
if portPos == -1 { if portPos == -1 {
hostHeader = fmt.Sprintf("Host: %s\n", rewriteHost) hostHeader = fmt.Sprintf("Host: %s\r\n", rewriteHost)
} else { } else {
hostHeader = fmt.Sprintf("Host: %s:%s\n", rewriteHost, kv[j+portPos+2:]) hostHeader = fmt.Sprintf("Host: %s:%s\r\n", rewriteHost, kv[j+portPos+2:])
} }
retBuf.WriteString(hostHeader) retBuf.WriteString(hostHeader)
peek = peek[i+1:] peek = peek[i+1:]