diff --git a/rafthttp/http.go b/rafthttp/http.go index 55a31f5d2..ea1b5fc6d 100644 --- a/rafthttp/http.go +++ b/rafthttp/http.go @@ -91,11 +91,7 @@ func (h *pipelineHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - if from, err := types.IDFromString(r.Header.Get("X-Server-From")); err == nil { - if urls := r.Header.Get("X-PeerURLs"); urls != "" { - h.tr.AddRemote(from, strings.Split(urls, ",")) - } - } + addRemoteFromRequest(h.tr, r) // Limit the data size that could be read from the request body, which ensures that read from // connection will not time out accidentally due to possible blocking in underlying implementation. @@ -176,11 +172,7 @@ func (h *snapshotHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - if from, err := types.IDFromString(r.Header.Get("X-Server-From")); err == nil { - if urls := r.Header.Get("X-PeerURLs"); urls != "" { - h.tr.AddRemote(from, strings.Split(urls, ",")) - } - } + addRemoteFromRequest(h.tr, r) dec := &messageDecoder{r: r.Body} // let snapshots be very large since they can exceed 512MB for large installations diff --git a/rafthttp/util.go b/rafthttp/util.go index 12e548c77..2f78415fe 100644 --- a/rafthttp/util.go +++ b/rafthttp/util.go @@ -175,3 +175,14 @@ func setPeerURLsHeader(req *http.Request, urls types.URLs) { } req.Header.Set("X-PeerURLs", strings.Join(peerURLs, ",")) } + +// addRemoteFromRequest add remote according to request header +func addRemoteFromRequest(tr Transporter, r *http.Request) bool { + if from, err := types.IDFromString(r.Header.Get("X-Server-From")); err == nil { + if urls := r.Header.Get("X-PeerURLs"); urls != "" { + tr.AddRemote(from, strings.Split(urls, ",")) + return true + } + } + return false +}