From 38c72221d6642a3cb7b81051e9577280e8868945 Mon Sep 17 00:00:00 2001 From: rick <1450685+LinuxSuRen@users.noreply.github.com> Date: Mon, 25 Oct 2021 16:05:28 +0800 Subject: [PATCH] Reject unreachbale candidate servers --- cmd/server/root.go | 1 + pkg/proxy/candidate_test.go | 32 ++++++++++++++++++--- pkg/proxy/handler.go | 13 +++++---- pkg/proxy/types.go | 56 +++++++++++++++++++++++++++++++++++++ pkg/server/health.go | 9 ++++++ 5 files changed, 102 insertions(+), 9 deletions(-) create mode 100644 pkg/server/health.go diff --git a/cmd/server/root.go b/cmd/server/root.go index b309ec5..5524be2 100644 --- a/cmd/server/root.go +++ b/cmd/server/root.go @@ -53,6 +53,7 @@ func createServerCommand() (cmd *cobra.Command) { func (o *serverOption) runE(cmd *cobra.Command, args []string) (err error) { switch o.mode { case "server": + http.HandleFunc("/health", server.HealthHandler) http.HandleFunc("/", server.GogetHandler) if err = server.IntervalSelfRegistry(o.proxyCenter, o.externalAddress, time.Minute*1); err != nil { err = fmt.Errorf("failed to self registry to the center proxy, error: %v", err) diff --git a/pkg/proxy/candidate_test.go b/pkg/proxy/candidate_test.go index f9deb31..5bf93d8 100644 --- a/pkg/proxy/candidate_test.go +++ b/pkg/proxy/candidate_test.go @@ -2,11 +2,12 @@ package proxy import ( "github.com/stretchr/testify/assert" + "sort" "testing" "time" ) -func TestCandidate(t *testing.T) { +func TestCandidateSlice(t *testing.T) { candidates := candidateSlice{ []candidate{{ address: "fake", @@ -30,6 +31,7 @@ func TestCandidate(t *testing.T) { assert.True(t, ok, "should be able to find the alive candidate which was updated") emptyCandidates := candidateSlice{} + assert.Nil(t, emptyCandidates.first()) emptyCandidates.addCandidate("fake") _, ok = emptyCandidates.findAlive() assert.True(t, ok, "should be able to find the alive candidate") @@ -39,7 +41,7 @@ func TestCandidate(t *testing.T) { expiredCandidates := candidateSlice{ candidates: []candidate{{ - address: "fake", + address: "fake", heartBeat: time.Now().Add(time.Minute * -10), }}, } @@ -60,7 +62,7 @@ func TestCandidatesHelper(t *testing.T) { // valid candidates array candidatesArray = []interface{}{ map[interface{}]interface{}{ - "address": "fake", + "address": "fake", "heartBeat": time.Now().Format(timeFormat), }, } @@ -73,7 +75,7 @@ func TestCandidatesHelper(t *testing.T) { // from map candidatesMap := []map[interface{}]interface{}{ { - "address": "fake", + "address": "fake", "heartBeat": time.Now(), }, } @@ -83,3 +85,25 @@ func TestCandidatesHelper(t *testing.T) { assert.True(t, ok) assert.Equal(t, "fake", aliveCandidate.address) } + +func TestCandidateSliceSort(t *testing.T) { + candidates := &candidateSlice{candidates: []candidate{{ + address: "one", + heartBeat: time.Now().Add(time.Minute), + }, { + address: "two", + heartBeat: time.Now().Add(time.Minute * 2), + }}} + sort.Sort(candidates) + firstCandidate := candidates.first() + assert.NotNil(t, firstCandidate) + assert.Equal(t, firstCandidate.address, "two") +} + +func TestCandidate(t *testing.T) { + candidate := NewCandidate("http://fake") + assert.Equal(t, "fake", candidate.getHost()) + + candidate = NewCandidate("https://fake") + assert.Equal(t, "fake", candidate.getHost()) +} diff --git a/pkg/proxy/handler.go b/pkg/proxy/handler.go index 58794b2..7b474fc 100644 --- a/pkg/proxy/handler.go +++ b/pkg/proxy/handler.go @@ -42,9 +42,12 @@ func isValid(uri string) bool { // RegistryHandler receive the proxy registry request func RegistryHandler(w http.ResponseWriter, r *http.Request) { - address := r.URL.Query().Get("address") - address = strings.ReplaceAll(address, "http://", "") - address = strings.ReplaceAll(address, "https://", "") + candidate := NewCandidate(r.URL.Query().Get("address")) + if !candidate.reachable() { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(fmt.Sprintf("%s is not reachable", candidate.address))) + return + } var ( candidates *candidateSlice @@ -59,9 +62,9 @@ func RegistryHandler(w http.ResponseWriter, r *http.Request) { candidates = newFromArray(candidatesRaw) } - candidates.addCandidate(address) + candidates.addCandidate(candidate.getHost()) - fmt.Println("receive candidate server", address) + fmt.Println("receive candidate server", candidate.getHost()) if err := saveCandidates(candidates); err == nil { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("ok")) diff --git a/pkg/proxy/types.go b/pkg/proxy/types.go index 0bb5398..6ba2459 100644 --- a/pkg/proxy/types.go +++ b/pkg/proxy/types.go @@ -2,6 +2,9 @@ package proxy import ( "fmt" + "net/http" + "sort" + "strings" "time" ) @@ -11,6 +14,36 @@ type candidate struct { expired bool } +// NewCandidate creates a new candidate instance +func NewCandidate(address string) *candidate { + return &candidate{address: address} +} + +func (c *candidate) reachable() bool { + if c.address == "" { + return false + } + + var address string + if strings.HasPrefix(c.address, "https://") || strings.HasPrefix(c.address, "http://") { + address = c.address + } else { + address = fmt.Sprintf("http://%s", c.address) + } + + client := http.DefaultClient + client.Timeout = time.Second * 3 + resp, err := client.Get(address) + return err == nil && resp.StatusCode == http.StatusOK +} + +func (c *candidate) getHost() (address string) { + address = c.address + address = strings.ReplaceAll(address, "http://", "") + address = strings.ReplaceAll(address, "https://", "") + return +} + var aliveDuration = time.Minute * 2 const timeFormat = time.RFC3339 @@ -19,7 +52,16 @@ type candidateSlice struct { candidates []candidate } +func (c *candidateSlice) first() *candidate { + if len(c.candidates) > 0 { + return &c.candidates[0] + } + return nil +} + func (c *candidateSlice) findAlive() (candidate, bool) { + sort.Sort(c) + for i, _ := range c.candidates { candidateItem := c.candidates[i] if candidateItem.address != "" && candidateItem.heartBeat.Add(aliveDuration).After(time.Now()) { @@ -59,6 +101,20 @@ func (c *candidateSlice) size() int { return len(c.candidates) } +func (c *candidateSlice) Len() int { + return c.size() +} + +func (c *candidateSlice) Less(i, j int) bool { + left := c.candidates[i].heartBeat + right := c.candidates[j].heartBeat + return left.After(right) +} + +func (c *candidateSlice) Swap(i, j int) { + c.candidates[i], c.candidates[j] = c.candidates[j], c.candidates[i] +} + func (c *candidateSlice) getMap() (result []map[interface{}]interface{}) { result = make([]map[interface{}]interface{}, 0) for i, _ := range c.candidates { diff --git a/pkg/server/health.go b/pkg/server/health.go new file mode 100644 index 0000000..dba8238 --- /dev/null +++ b/pkg/server/health.go @@ -0,0 +1,9 @@ +package server + +import "net/http" + +// HealthHandler is the handler of the server health request +func HealthHandler(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) +}