4c96a297b9f0ccc0923a19448cadbe6c11570b9c

Build # 4c96a297b9f0ccc0923a19448cadbe6c11570b9c

Build Type

push

gitlab-ci

Committed by GitHub

Commit Message

Merge pull request #1209 from merlin-northern/men_5915_lat_lon_in_inv_geo

Run Details

9806 of 12349 relevant lines covered (79.41%)

30.04 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.64

/app/proxy/proxy_ws.go

// Copyright 2022 Northern.tech AS
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
package proxy

// Inspired by https://github.com/koding/websocketproxy

import (
        "context"
        "fmt"
        "net/http"
        "net/url"
        "strings"
        "sync"

        "github.com/gorilla/websocket"
)

const (
        maxWsConnections = 1
)

func newUpgrader() *websocket.Upgrader {
        return &websocket.Upgrader{
                CheckOrigin: func(r *http.Request) bool {
                        return true
                },
        }
}

func wsSchemeFromHttpScheme(httpScheme string) string {
        switch httpScheme {
        case "http":
                return "ws"
        case "https":
                return "wss"
        default:
                return ""
        }
}

func (pc *ProxyController) wsRunning() bool {
        pc.wsConnectionsMutex.Lock()
        defer pc.wsConnectionsMutex.Unlock()
        return len(pc.wsConnections) > 0
}

func (pc *proxyControllerInner) wsAvailable() bool {
        pc.wsConnectionsMutex.Lock()
        defer pc.wsConnectionsMutex.Unlock()
        return len(pc.wsConnections) < maxWsConnections
}

func (pc *proxyControllerInner) DoWsUpgrade(
        ctx context.Context,
        w http.ResponseWriter,
        r *http.Request) {
        // Convert server request to client request and override r.URL
        r.RequestURI = ""
        r.Host = ""
        r.URL.Scheme = wsSchemeFromHttpScheme(pc.conf.backend.Scheme)
        r.URL.Host = pc.conf.backend.Host

        // Copy all headers but delete the websocket's handshake related ones
        requestHeader := http.Header{}
        copyHeader(requestHeader, r.Header)
        requestHeader.Del("Sec-Websocket-Key")
        requestHeader.Del("Sec-Websocket-Version")
        requestHeader.Del("Upgrade")
        requestHeader.Del("Connection")

        wsUrl := url.URL{
                Host:   pc.conf.backend.Host,
                Scheme: wsSchemeFromHttpScheme(pc.conf.backend.Scheme),
                Path:   ApiUrlDevicesConnect,
        }

        connBackend, resp, err := pc.wsDialer.DialContext(ctx, wsUrl.String(), requestHeader)
        if err != nil {
                log.Errorf("couldn't dial to remote backend url %q, err: %s", wsUrl.String(), err.Error())
                if resp != nil {
                        // WebSocket handshake failed, reply the client with backend's resp
                        if err := copyResponse(w, resp); err != nil {
                                log.Errorf("couldn't write response after failed remote backend handshake: %s", err)
                        }
                } else {
                        http.Error(
                                w,
                                http.StatusText(http.StatusServiceUnavailable),
                                http.StatusServiceUnavailable,
                        )
                }
                return
        }
        defer connBackend.Close()

        // Only pass certain headers to the upgrader.
        upgradeHeader := http.Header{}
        if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {
                upgradeHeader.Set("Sec-Websocket-Protocol", hdr)
        }
        if hdr := resp.Header.Get("Set-Cookie"); hdr != "" {
                upgradeHeader.Set("Set-Cookie", hdr)
        }

        // If backend replied with a protocol, use that one for our upgrader
        // It shall only be one after the handshake, but Upgrader.Subprotocols
        // expects an slice, so copy "all".
        upgrader := newUpgrader()
        if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {
                backendProtocols := strings.Split(hdr, ",")
                for i := range backendProtocols {
                        backendProtocols[i] = strings.TrimSpace(backendProtocols[i])
                }
                upgrader.Subprotocols = backendProtocols
        }

        // Now ready to upgrade connection from client to proxy!
        connClient, err := upgrader.Upgrade(w, r, upgradeHeader)
        if err != nil {
                log.Errorf("couldn't upgrade %s", err)
                return
        }
        defer connClient.Close()

        newConn := &wsConnection{
                connBackend: connBackend,
                connClient:  connClient,
        }
        pc.wsConnectionsMutex.Lock()
        pc.wsConnections[newConn] = true
        pc.wsConnectionsMutex.Unlock()

        errClient := make(chan error, 1)
        errBackend := make(chan error, 1)

        go pc.forwardWsConnection(
                newConn.connClient,
                newConn.connBackend,
                errClient,
                &newConn.connClientWriteMutex,
        )
        go pc.forwardWsConnection(
                newConn.connBackend,
                newConn.connClient,
                errBackend,
                &newConn.connBackendWriteMutex,
        )

        var messageF string
        select {
        case err = <-errClient:
                messageF = "error forwarding from backend to client: %v"
        case err = <-errBackend:
                messageF = "error forwarding from client to backend: %v"

        }
        if e, ok := err.(*websocket.CloseError); !ok || e.Code != websocket.CloseNormalClosure {
                log.Errorf(messageF, err.Error())
        }

        pc.wsConnectionsMutex.Lock()
        delete(pc.wsConnections, newConn)
        pc.wsConnectionsMutex.Unlock()
}

func (pc *proxyControllerInner) forwardWsConnection(
        dst, src *websocket.Conn,
        errChann chan error,
        mutex *sync.Mutex,
) {
        for {
                msgType, msg, err := src.ReadMessage()
                if err != nil {
                        m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err))
                        if e, ok := err.(*websocket.CloseError); ok {
                                if e.Code != websocket.CloseNoStatusReceived {
                                        m = websocket.FormatCloseMessage(e.Code, e.Text)
                                }
                        }
                        mutex.Lock()
                        errWrite := dst.WriteMessage(websocket.CloseMessage, m)
                        mutex.Unlock()
                        if errWrite != nil {
                                log.Warningf("error while sending close message: %v", errWrite.Error())
                        }
                        errChann <- err
                        break
                }
                mutex.Lock()
                err = dst.WriteMessage(msgType, msg)
                mutex.Unlock()
                if err != nil {
                        errChann <- err
                        break
                }
        }
}

func (pc *ProxyController) CloseWsConnections() {
        log.Info("shutting down websocket connections")
        m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "shutting down proxy")

        pc.wsConnectionsMutex.Lock()
        defer pc.wsConnectionsMutex.Unlock()
        for c := range pc.wsConnections {
                c.connBackendWriteMutex.Lock()
                errWrite := c.connBackend.WriteMessage(websocket.CloseMessage, m)
                if errWrite != nil {
                        log.Errorf("error while sending close message to backend: %v", errWrite.Error())
                }
                c.connBackendWriteMutex.Unlock()

                c.connClientWriteMutex.Lock()
                errWrite = c.connClient.WriteMessage(websocket.CloseMessage, m)
                if errWrite != nil {
                        log.Errorf("error while sending close message to client: %v", errWrite.Error())
                }
                c.connClientWriteMutex.Unlock()

                delete(pc.wsConnections, c)
        }
}

1	// Copyright 2022 Northern.tech AS
2	//
3	// Licensed under the Apache License, Version 2.0 (the "License");
4	// you may not use this file except in compliance with the License.
5	// You may obtain a copy of the License at
6	//
7	// http://www.apache.org/licenses/LICENSE-2.0
8	//
9	// Unless required by applicable law or agreed to in writing, software
10	// distributed under the License is distributed on an "AS IS" BASIS,
11	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	// See the License for the specific language governing permissions and
13	// limitations under the License.
14	package proxy
15
16	// Inspired by https://github.com/koding/websocketproxy
17
18	import (
19	"context"
20	"fmt"
21	"net/http"
22	"net/url"
23	"strings"
24	"sync"
25
26	"github.com/gorilla/websocket"
27	)
28
29	const (
30	maxWsConnections = 1
31	)
32
33	func newUpgrader() *websocket.Upgrader {	2✔
34	return &websocket.Upgrader{	2✔
35	CheckOrigin: func(r *http.Request) bool {	4✔
36	return true	2✔
37	},	2✔
38	}
39	}
40
41	func wsSchemeFromHttpScheme(httpScheme string) string {	2✔
42	switch httpScheme {	2✔
43	case "http":	2✔
44	return "ws"	2✔
45	case "https":	2✔
46	return "wss"	2✔
47	default:	×
48	return ""	×
49	}
50	}
51
52	func (pc *ProxyController) wsRunning() bool {	17✔
53	pc.wsConnectionsMutex.Lock()	17✔
54	defer pc.wsConnectionsMutex.Unlock()	17✔
55	return len(pc.wsConnections) > 0	17✔
56	}	17✔
57
58	func (pc *proxyControllerInner) wsAvailable() bool {	2✔
59	pc.wsConnectionsMutex.Lock()	2✔
60	defer pc.wsConnectionsMutex.Unlock()	2✔
61	return len(pc.wsConnections) < maxWsConnections	2✔
62	}	2✔
63
64	func (pc *proxyControllerInner) DoWsUpgrade(
65	ctx context.Context,
66	w http.ResponseWriter,
67	r *http.Request) {	2✔
68	// Convert server request to client request and override r.URL	2✔
69	r.RequestURI = ""	2✔
70	r.Host = ""	2✔
71	r.URL.Scheme = wsSchemeFromHttpScheme(pc.conf.backend.Scheme)	2✔
72	r.URL.Host = pc.conf.backend.Host	2✔
73		2✔
74	// Copy all headers but delete the websocket's handshake related ones	2✔
75	requestHeader := http.Header{}	2✔
76	copyHeader(requestHeader, r.Header)	2✔
77	requestHeader.Del("Sec-Websocket-Key")	2✔
78	requestHeader.Del("Sec-Websocket-Version")	2✔
79	requestHeader.Del("Upgrade")	2✔
80	requestHeader.Del("Connection")	2✔
81		2✔
82	wsUrl := url.URL{	2✔
83	Host: pc.conf.backend.Host,	2✔
84	Scheme: wsSchemeFromHttpScheme(pc.conf.backend.Scheme),	2✔
85	Path: ApiUrlDevicesConnect,	2✔
86	}	2✔
87		2✔
88	connBackend, resp, err := pc.wsDialer.DialContext(ctx, wsUrl.String(), requestHeader)	2✔
89	if err != nil {	2✔
90	log.Errorf("couldn't dial to remote backend url %q, err: %s", wsUrl.String(), err.Error())	×
91	if resp != nil {	×
92	// WebSocket handshake failed, reply the client with backend's resp	×
93	if err := copyResponse(w, resp); err != nil {	×
94	log.Errorf("couldn't write response after failed remote backend handshake: %s", err)	×
95	}	×
96	} else {	×
97	http.Error(	×
98	w,	×
99	http.StatusText(http.StatusServiceUnavailable),	×
100	http.StatusServiceUnavailable,	×
101	)	×
102	}	×
103	return	×
104	}
105	defer connBackend.Close()	2✔
106		2✔
107	// Only pass certain headers to the upgrader.	2✔
108	upgradeHeader := http.Header{}	2✔
109	if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {	4✔
110	upgradeHeader.Set("Sec-Websocket-Protocol", hdr)	2✔
111	}	2✔
112	if hdr := resp.Header.Get("Set-Cookie"); hdr != "" {	2✔
113	upgradeHeader.Set("Set-Cookie", hdr)	×
114	}	×
115
116	// If backend replied with a protocol, use that one for our upgrader
117	// It shall only be one after the handshake, but Upgrader.Subprotocols
118	// expects an slice, so copy "all".
119	upgrader := newUpgrader()	2✔
120	if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {	4✔
121	backendProtocols := strings.Split(hdr, ",")	2✔
122	for i := range backendProtocols {	4✔
123	backendProtocols[i] = strings.TrimSpace(backendProtocols[i])	2✔
124	}	2✔
125	upgrader.Subprotocols = backendProtocols	2✔
126	}
127
128	// Now ready to upgrade connection from client to proxy!
129	connClient, err := upgrader.Upgrade(w, r, upgradeHeader)	2✔
130	if err != nil {	2✔
131	log.Errorf("couldn't upgrade %s", err)	×
132	return	×
133	}	×
134	defer connClient.Close()	2✔
135		2✔
136	newConn := &wsConnection{	2✔
137	connBackend: connBackend,	2✔
138	connClient: connClient,	2✔
139	}	2✔
140	pc.wsConnectionsMutex.Lock()	2✔
141	pc.wsConnections[newConn] = true	2✔
142	pc.wsConnectionsMutex.Unlock()	2✔
143		2✔
144	errClient := make(chan error, 1)	2✔
145	errBackend := make(chan error, 1)	2✔
146		2✔
147	go pc.forwardWsConnection(	2✔
148	newConn.connClient,	2✔
149	newConn.connBackend,	2✔
150	errClient,	2✔
151	&newConn.connClientWriteMutex,	2✔
152	)	2✔
153	go pc.forwardWsConnection(	2✔
154	newConn.connBackend,	2✔
155	newConn.connClient,	2✔
156	errBackend,	2✔
157	&newConn.connBackendWriteMutex,	2✔
158	)	2✔
159		2✔
160	var messageF string	2✔
161	select {	2✔
162	case err = <-errClient:	2✔
163	messageF = "error forwarding from backend to client: %v"	2✔
164	case err = <-errBackend:	2✔
165	messageF = "error forwarding from client to backend: %v"	2✔
166
167	}
168	if e, ok := err.(*websocket.CloseError); !ok \|\| e.Code != websocket.CloseNormalClosure {	4✔
169	log.Errorf(messageF, err.Error())	2✔
170	}	2✔
171
172	pc.wsConnectionsMutex.Lock()	2✔
173	delete(pc.wsConnections, newConn)	2✔
174	pc.wsConnectionsMutex.Unlock()	2✔
175	}
176
177	func (pc *proxyControllerInner) forwardWsConnection(
178	dst, src *websocket.Conn,
179	errChann chan error,
180	mutex *sync.Mutex,
181	) {	2✔
182	for {	4✔
183	msgType, msg, err := src.ReadMessage()	2✔
184	if err != nil {	4✔
185	m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err))	2✔
186	if e, ok := err.(*websocket.CloseError); ok {	4✔
187	if e.Code != websocket.CloseNoStatusReceived {	4✔
188	m = websocket.FormatCloseMessage(e.Code, e.Text)	2✔
189	}	2✔
190	}
191	mutex.Lock()	2✔
192	errWrite := dst.WriteMessage(websocket.CloseMessage, m)	2✔
193	mutex.Unlock()	2✔
194	if errWrite != nil {	4✔
195	log.Warningf("error while sending close message: %v", errWrite.Error())	2✔
196	}	2✔
197	errChann <- err	2✔
198	break	2✔
199	}
200	mutex.Lock()	2✔
201	err = dst.WriteMessage(msgType, msg)	2✔
202	mutex.Unlock()	2✔
203	if err != nil {	2✔
204	errChann <- err	×
205	break	×
206	}
207	}
208	}
209
210	func (pc *ProxyController) CloseWsConnections() {	2✔
211	log.Info("shutting down websocket connections")	2✔
212	m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "shutting down proxy")	2✔
213		2✔
214	pc.wsConnectionsMutex.Lock()	2✔
215	defer pc.wsConnectionsMutex.Unlock()	2✔
216	for c := range pc.wsConnections {	4✔
217	c.connBackendWriteMutex.Lock()	2✔
218	errWrite := c.connBackend.WriteMessage(websocket.CloseMessage, m)	2✔
219	if errWrite != nil {	4✔
220	log.Errorf("error while sending close message to backend: %v", errWrite.Error())	2✔
221	}	2✔
222	c.connBackendWriteMutex.Unlock()	2✔
223		2✔
224	c.connClientWriteMutex.Lock()	2✔
225	errWrite = c.connClient.WriteMessage(websocket.CloseMessage, m)	2✔
226	if errWrite != nil {	2✔
227	log.Errorf("error while sending close message to client: %v", errWrite.Error())	×
228	}	×
229	c.connClientWriteMutex.Unlock()	2✔
230		2✔
231	delete(pc.wsConnections, c)	2✔
232	}
233	}

mendersoftware / mender / 4c96a297b9f0ccc0923a19448cadbe6c11570b9c

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous