Merge pull request #14 from LinuxSuRen/server-reachable-check

Reject unreachbale candidate servers

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)

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())
+}

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"))

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 {

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"))
+}