Sfoglia il codice sorgente

Merge pull request #702 from unihorn/60

feat(server): make header-only requests work
Yicheng Qin 11 anni fa
parent
commit
4ce8c3499b

+ 29 - 17
server/server.go

@@ -100,33 +100,33 @@ func (s *Server) Store() store.Store {
 }
 
 func (s *Server) installV1(r *mux.Router) {
-	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.GetKeyHandler).Methods("GET")
+	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.GetKeyHandler).Methods("GET", "HEAD")
 	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.SetKeyHandler).Methods("POST", "PUT")
 	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.DeleteKeyHandler).Methods("DELETE")
-	s.handleFuncV1(r, "/v1/watch/{key:.*}", v1.WatchKeyHandler).Methods("GET", "POST")
-	s.handleFunc(r, "/v1/leader", s.GetLeaderHandler).Methods("GET")
-	s.handleFunc(r, "/v1/machines", s.GetPeersHandler).Methods("GET")
-	s.handleFunc(r, "/v1/peers", s.GetPeersHandler).Methods("GET")
-	s.handleFunc(r, "/v1/stats/self", s.GetStatsHandler).Methods("GET")
-	s.handleFunc(r, "/v1/stats/leader", s.GetLeaderStatsHandler).Methods("GET")
-	s.handleFunc(r, "/v1/stats/store", s.GetStoreStatsHandler).Methods("GET")
+	s.handleFuncV1(r, "/v1/watch/{key:.*}", v1.WatchKeyHandler).Methods("GET", "HEAD", "POST")
+	s.handleFunc(r, "/v1/leader", s.GetLeaderHandler).Methods("GET", "HEAD")
+	s.handleFunc(r, "/v1/machines", s.GetPeersHandler).Methods("GET", "HEAD")
+	s.handleFunc(r, "/v1/peers", s.GetPeersHandler).Methods("GET", "HEAD")
+	s.handleFunc(r, "/v1/stats/self", s.GetStatsHandler).Methods("GET", "HEAD")
+	s.handleFunc(r, "/v1/stats/leader", s.GetLeaderStatsHandler).Methods("GET", "HEAD")
+	s.handleFunc(r, "/v1/stats/store", s.GetStoreStatsHandler).Methods("GET", "HEAD")
 }
 
 func (s *Server) installV2(r *mux.Router) {
 	r2 := mux.NewRouter()
 	r.PathPrefix("/v2").Handler(ehttp.NewLowerQueryParamsHandler(r2))
 
-	s.handleFuncV2(r2, "/v2/keys/{key:.*}", v2.GetHandler).Methods("GET")
+	s.handleFuncV2(r2, "/v2/keys/{key:.*}", v2.GetHandler).Methods("GET", "HEAD")
 	s.handleFuncV2(r2, "/v2/keys/{key:.*}", v2.PostHandler).Methods("POST")
 	s.handleFuncV2(r2, "/v2/keys/{key:.*}", v2.PutHandler).Methods("PUT")
 	s.handleFuncV2(r2, "/v2/keys/{key:.*}", v2.DeleteHandler).Methods("DELETE")
-	s.handleFunc(r2, "/v2/leader", s.GetLeaderHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/machines", s.GetPeersHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/peers", s.GetPeersHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/stats/self", s.GetStatsHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/stats/leader", s.GetLeaderStatsHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/stats/store", s.GetStoreStatsHandler).Methods("GET")
-	s.handleFunc(r2, "/v2/speedTest", s.SpeedTestHandler).Methods("GET")
+	s.handleFunc(r2, "/v2/leader", s.GetLeaderHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/machines", s.GetPeersHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/peers", s.GetPeersHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/stats/self", s.GetStatsHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/stats/leader", s.GetLeaderStatsHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/stats/store", s.GetStoreStatsHandler).Methods("GET", "HEAD")
+	s.handleFunc(r2, "/v2/speedTest", s.SpeedTestHandler).Methods("GET", "HEAD")
 }
 
 func (s *Server) installMod(r *mux.Router) {
@@ -134,7 +134,7 @@ func (s *Server) installMod(r *mux.Router) {
 }
 
 func (s *Server) installDebug(r *mux.Router) {
-	s.handleFunc(r, "/debug/metrics", s.GetMetricsHandler).Methods("GET")
+	s.handleFunc(r, "/debug/metrics", s.GetMetricsHandler).Methods("GET", "HEAD")
 	r.HandleFunc("/debug/pprof", pprof.Index)
 	r.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
 	r.HandleFunc("/debug/pprof/profile", pprof.Profile)
@@ -156,11 +156,23 @@ func (s *Server) handleFuncV2(r *mux.Router, path string, f func(http.ResponseWr
 	})
 }
 
+type HEADResponseWriter struct {
+	http.ResponseWriter
+}
+
+func (w *HEADResponseWriter) Write([]byte) (int, error) {
+	return 0, nil
+}
+
 // Adds a server handler to the router.
 func (s *Server) handleFunc(r *mux.Router, path string, f func(http.ResponseWriter, *http.Request) error) *mux.Route {
 
 	// Wrap the standard HandleFunc interface to pass in the server reference.
 	return r.HandleFunc(path, func(w http.ResponseWriter, req *http.Request) {
+		if req.Method == "HEAD" {
+			w = &HEADResponseWriter{w}
+		}
+
 		// Log request.
 		log.Debugf("[recv] %s %s %s [%s]", req.Method, s.URL(), req.URL.Path, req.RemoteAddr)
 

+ 6 - 3
server/v1/get_key_handler.go

@@ -18,11 +18,14 @@ func GetKeyHandler(w http.ResponseWriter, req *http.Request, s Server) error {
 		return err
 	}
 
+	w.WriteHeader(http.StatusOK)
+
+	if req.Method == "HEAD" {
+		return nil
+	}
+
 	// Convert event to a response and write to client.
 	b, _ := json.Marshal(event.Response(s.Store().Index()))
-
-	w.WriteHeader(http.StatusOK)
 	w.Write(b)
-
 	return nil
 }

+ 24 - 0
server/v1/tests/get_handler_test.go

@@ -176,3 +176,27 @@ func TestV1WatchKeyWithIndex(t *testing.T) {
 		assert.Equal(t, body["index"], 3, "")
 	})
 }
+
+// Ensures that HEAD works.
+//
+//   $ curl -I localhost:4001/v1/keys/foo/bar -> fail
+//   $ curl -X PUT localhost:4001/v1/keys/foo/bar -d value=XXX
+//   $ curl -I localhost:4001/v1/keys/foo/bar
+//
+func TestV1HeadKey(t *testing.T) {
+	tests.RunServer(func(s *server.Server) {
+		v := url.Values{}
+		v.Set("value", "XXX")
+		fullURL := fmt.Sprintf("%s%s", s.URL(), "/v1/keys/foo/bar")
+		resp, _ := tests.Get(fullURL)
+		assert.Equal(t, resp.StatusCode, http.StatusNotFound)
+		assert.Equal(t, resp.ContentLength, -1)
+
+		resp, _ = tests.PutForm(fullURL, v)
+		tests.ReadBody(resp)
+
+		resp, _ = tests.Get(fullURL)
+		assert.Equal(t, resp.StatusCode, http.StatusOK)
+		assert.Equal(t, resp.ContentLength, -1)
+	})
+}

+ 4 - 2
server/v1/watch_key_handler.go

@@ -32,9 +32,11 @@ func WatchKeyHandler(w http.ResponseWriter, req *http.Request, s Server) error {
 	event := <-watcher.EventChan
 
 	// Convert event to a response and write to client.
-	b, _ := json.Marshal(event.Response(s.Store().Index()))
 	w.WriteHeader(http.StatusOK)
+	if req.Method == "HEAD" {
+		return nil
+	}
+	b, _ := json.Marshal(event.Response(s.Store().Index()))
 	w.Write(b)
-
 	return nil
 }

+ 14 - 4
server/v2/get_handler.go

@@ -41,13 +41,13 @@ func GetHandler(w http.ResponseWriter, req *http.Request, s Server) error {
 	stream := (req.FormValue("stream") == "true")
 
 	if req.FormValue("wait") == "true" {
-		return handleWatch(key, recursive, stream, waitIndex, w, s)
+		return handleWatch(key, recursive, stream, waitIndex, w, req, s)
 	}
 
-	return handleGet(key, recursive, sort, w, s)
+	return handleGet(key, recursive, sort, w, req, s)
 }
 
-func handleWatch(key string, recursive, stream bool, waitIndex string, w http.ResponseWriter, s Server) error {
+func handleWatch(key string, recursive, stream bool, waitIndex string, w http.ResponseWriter, req *http.Request, s Server) error {
 	// Create a command to watch from a given index (default 0).
 	var sinceIndex uint64 = 0
 	var err error
@@ -84,6 +84,9 @@ func handleWatch(key string, recursive, stream bool, waitIndex string, w http.Re
 					// send to the client in time. Then we simply end streaming.
 					return nil
 				}
+				if req.Method == "HEAD" {
+					continue
+				}
 
 				b, _ := json.Marshal(event)
 				_, err := w.Write(b)
@@ -99,18 +102,25 @@ func handleWatch(key string, recursive, stream bool, waitIndex string, w http.Re
 	case <-closeChan:
 		watcher.Remove()
 	case event := <-watcher.EventChan:
+		if req.Method == "HEAD" {
+			return nil
+		}
 		b, _ := json.Marshal(event)
 		w.Write(b)
 	}
 	return nil
 }
 
-func handleGet(key string, recursive, sort bool, w http.ResponseWriter, s Server) error {
+func handleGet(key string, recursive, sort bool, w http.ResponseWriter, req *http.Request, s Server) error {
 	event, err := s.Store().Get(key, recursive, sort)
 	if err != nil {
 		return err
 	}
 
+	if req.Method == "HEAD" {
+		return nil
+	}
+
 	writeHeaders(w, s)
 	b, _ := json.Marshal(event)
 	w.Write(b)

+ 24 - 0
server/v2/tests/get_handler_test.go

@@ -229,3 +229,27 @@ func TestV2WatchKeyInDir(t *testing.T) {
 		assert.Equal(t, node["key"], "/keyindir", "")
 	})
 }
+
+// Ensures that HEAD could work.
+//
+//   $ curl -I localhost:4001/v2/keys/foo/bar -> fail
+//   $ curl -X PUT localhost:4001/v2/keys/foo/bar -d value=XXX
+//   $ curl -I localhost:4001/v2/keys/foo/bar
+//
+func TestV2HeadKey(t *testing.T) {
+	tests.RunServer(func(s *server.Server) {
+		v := url.Values{}
+		v.Set("value", "XXX")
+		fullURL := fmt.Sprintf("%s%s", s.URL(), "/v2/keys/foo/bar")
+		resp, _ := tests.Head(fullURL)
+		assert.Equal(t, resp.StatusCode, http.StatusNotFound)
+		assert.Equal(t, resp.ContentLength, -1)
+
+		resp, _ = tests.PutForm(fullURL, v)
+		tests.ReadBody(resp)
+
+		resp, _ = tests.Head(fullURL)
+		assert.Equal(t, resp.StatusCode, http.StatusOK)
+		assert.Equal(t, resp.ContentLength, -1)
+	})
+}

+ 4 - 0
tests/http_utils.go

@@ -35,6 +35,10 @@ func ReadBodyJSON(resp *http.Response) map[string]interface{} {
 	return m
 }
 
+func Head(url string) (*http.Response, error) {
+	return send("HEAD", url, "application/json", nil)
+}
+
 func Get(url string) (*http.Response, error) {
 	return send("GET", url, "application/json", nil)
 }