mosn

1
package util
2

3
import (
4
	"fmt"
5
	"net"
6
	"net/http"
7
	"net/http/httptest"
8
	"strings"
9
	"sync"
10
	"testing"
11
	"time"
12

13
	"golang.org/x/net/http2"
14
)
15

16
type UpstreamServer interface {
17
	GoServe()
18
	Close()
19
	Addr() string
20
}
21

22
type ServeConn func(t *testing.T, conn net.Conn)
23
type upstreamServer struct {
24
	Listener net.Listener
25
	Serve    ServeConn
26
	Address  string
27
	conns    []net.Conn
28
	mu       sync.Mutex
29
	t        *testing.T
30
	closed   bool
31
}
32

33
func NewUpstreamServer(t *testing.T, addr string, serve ServeConn) UpstreamServer {
34
	return &upstreamServer{
35
		Serve:   serve,
36
		Address: addr,
37
		conns:   []net.Conn{},
38
		mu:      sync.Mutex{},
39
		t:       t,
40
	}
41
}
42

43
// GoServe and Close exists concurrency
44
// expected action is Close should close the server, should not restart
45
func (s *upstreamServer) GoServe() {
46
	s.mu.Lock()
47
	defer s.mu.Unlock()
48
	if s.closed {
49
		return
50
	}
51
	//try listen
52
	var err error
53
	for i := 0; i < 3; i++ {
54
		s.Listener, err = net.Listen("tcp", s.Address)
55
		if s.Listener != nil {
56
			break
57
		}
58
		time.Sleep(time.Second)
59
	}
60
	if s.Listener == nil {
61
		s.t.Fatalf("listen %s failed, error : %v\n", s.Address, err)
62
	}
63
	go s.serve()
64
}
65
func (s *upstreamServer) serve() {
66
	for {
67
		conn, err := s.Listener.Accept()
68
		if err != nil {
69
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
70
				s.t.Logf("Accept error: %v\n", err)
71
				continue
72
			}
73
			return
74
		}
75
		s.t.Logf("server %s Accept connection: %s\n", s.Listener.Addr().String(), conn.RemoteAddr().String())
76
		s.mu.Lock()
77
		if !s.closed {
78
			s.conns = append(s.conns, conn)
79
			go s.Serve(s.t, conn)
80
		}
81
		s.mu.Unlock()
82
	}
83
}
84

85
func (s *upstreamServer) Close() {
86
	s.mu.Lock()
87
	defer s.mu.Unlock()
88
	s.t.Logf("server %s closed\n", s.Address)
89
	if s.Listener != nil { // Close maybe called before GoServe
90
		s.Listener.Close()
91
	}
92
	for _, conn := range s.conns {
93
		conn.Close()
94
	}
95
	s.closed = true
96
}
97
func (s *upstreamServer) Addr() string {
98
	return s.Address
99
}
100

101
// Server Implement
102
type HTTPHandler struct{}
103

104
func (h *HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
105
	w.Header().Set("Content-Type", "text/plain")
106

107
	fmt.Fprintf(w, "\nRequestId:%s\n", r.Header.Get("Requestid"))
108

109
}
110

111
type HTTP2Server struct {
112
	t       *testing.T
113
	Server  *http2.Server
114
	Handler http.Handler
115
}
116

117
func (s *HTTP2Server) ServeConn(t *testing.T, conn net.Conn) {
118
	s.Server.ServeConn(conn, &http2.ServeConnOpts{Handler: s.Handler})
119
}
120

121
func NewUpstreamHTTP2(t *testing.T, addr string, h http.Handler) UpstreamServer {
122
	if h == nil {
123
		h = &HTTPHandler{}
124
	}
125
	s := &HTTP2Server{
126
		t:       t,
127
		Server:  &http2.Server{IdleTimeout: 1 * time.Minute},
128
		Handler: h,
129
	}
130
	return NewUpstreamServer(t, addr, s.ServeConn)
131
}
132

133
// Http Server
134
type HTTPServer struct {
135
	t       *testing.T
136
	server  *httptest.Server
137
	Handler http.Handler
138
}
139

140
func (s *HTTPServer) GoServe() {
141
	s.server.Start()
142
}
143
func (s *HTTPServer) Close() {
144
	s.server.Close()
145
}
146
func (s *HTTPServer) Addr() string {
147
	addr := strings.Split(s.server.URL, "http://")
148
	if len(addr) == 2 {
149
		return addr[1]
150
	}
151
	return ""
152
}
153

154
func NewHTTPServer(t *testing.T, h http.Handler) UpstreamServer {
155
	if h == nil {
156
		h = &HTTPHandler{}
157
	}
158
	s := &HTTPServer{t: t, Handler: h}
159
	s.server = httptest.NewUnstartedServer(s.Handler)
160
	return s
161
}
162