-
Notifications
You must be signed in to change notification settings - Fork 488
/
wrapper.go
120 lines (102 loc) · 4.17 KB
/
wrapper.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package api
import (
"fmt"
"io/ioutil"
"net/http"
"strings"
"time"
"build-booster/common/blog"
selfMetric "build-booster/server/pkg/metric"
"github.com/emicklei/go-restful"
)
const (
// HeaderRemote define a http-header name which store the true source ip from requester.
HeaderRemote = "source_ip"
)
// MasterRequired wrap the api handler and make it only available for master node.
func MasterRequired(f restful.RouteFunction) func(req *restful.Request, resp *restful.Response) {
return process(f, ProcessMasterOnly)
}
// NoLimit wrap the api handler and will process all requests.
func NoLimit(f restful.RouteFunction) func(req *restful.Request, resp *restful.Response) {
return process(f, ProcessNoLimit)
}
// Process log before and after a request. If options is mater-required, then redirect the request to master node and
// return the data from master node.
func process(f restful.RouteFunction, opts ProcessType) func(req *restful.Request, resp *restful.Response) {
return func(req *restful.Request, resp *restful.Response) {
entranceTime := time.Now().Local()
// handle the remote addr from other server for redirect
remote := strings.Split(req.Request.RemoteAddr, ":")
if len(remote) > 0 {
//if ok, _ := GetAPIResource().Rd.IsServerIP(remote[0]); !ok {
// req.Request.Header.Del(HeaderRemote)
// req.Request.Header.Set(HeaderRemote, req.Request.RemoteAddr)
//}
// TODO: trust the header remote ip for now, go back to check when rd is stable.
if req.Request.Header.Get(HeaderRemote) == "" {
req.Request.Header.Set(HeaderRemote, req.Request.RemoteAddr)
}
}
blog.Infof("Receive %s %s?%s From %s", req.Request.Method, req.Request.URL.Path, req.Request.URL.RawQuery, req.Request.Header.Get(HeaderRemote))
switch opts {
case ProcessNoLimit:
f(req, resp)
case ProcessMasterOnly:
isMaster, leader, err := GetAPIResource().Rd.IsMaster()
if err != nil {
blog.Errorf("process get master failed url(%s %s?%s): %v",
req.Request.Method, req.Request.URL.Path, req.Request.URL.RawQuery, err)
ReturnRest(&RestResponse{Resp: resp, ErrCode: ServerErrPreProcessFailed, HTTPCode: http.StatusInternalServerError})
} else if isMaster {
f(req, resp)
} else {
redirect(leader.GetURI(), req, resp)
}
}
useTime := time.Since(entranceTime).Nanoseconds() / 1000 / 1000
// record request metrics data
selfMetric.HttpRequestController.Observe(selfMetric.HttpRequestLabel{
Code: fmt.Sprintf("%d", resp.StatusCode()),
Method: req.Request.Method,
Path: req.Request.URL.Path,
}, float64(useTime))
blog.Infof("Return [%d] %dms %s %s?%s To %s", resp.StatusCode(), useTime, req.Request.Method, req.Request.URL.Path, req.Request.URL.RawQuery, req.Request.Header.Get(HeaderRemote))
}
}
type ProcessType string
const (
ProcessMasterOnly ProcessType = "master_only"
ProcessNoLimit ProcessType = "no_limit"
)
// redirect is like a proxy, it requests to the other node and return the data from that one.
func redirect(uri string, req *restful.Request, resp *restful.Response) {
uri += req.Request.URL.Path + "?" + req.Request.URL.RawQuery
blog.Infof("redirect to uri: %s %s", req.Request.Method, uri)
r, err := http.NewRequest(req.Request.Method, uri, req.Request.Body)
if err != nil {
blog.Errorf("redirect to uri(%s %s), get new request failed: %v", req.Request.Method, uri, err)
ReturnRest(&RestResponse{Resp: resp, ErrCode: ServerErrRedirectFailed, HTTPCode: http.StatusInternalServerError})
return
}
r.Header = req.Request.Header
r.Close = true
rsp, err := defaultClient.Do(r)
if err != nil {
blog.Errorf("redirect to uri(%s %s), do request failed: %v", req.Request.Method, uri, err)
ReturnRest(&RestResponse{Resp: resp, ErrCode: ServerErrRedirectFailed, HTTPCode: http.StatusInternalServerError})
return
}
defer func() {
_ = rsp.Body.Close()
}()
data, err := ioutil.ReadAll(rsp.Body)
if err != nil {
blog.Errorf("redirect to uri(%s %s) failed: %v", req.Request.Method, uri, err)
ReturnRest(&RestResponse{Resp: resp, ErrCode: ServerErrRedirectFailed, HTTPCode: http.StatusInternalServerError})
return
}
resp.WriteHeader(rsp.StatusCode)
_, _ = resp.ResponseWriter.Write(data)
}
var defaultClient = &http.Client{}