File: conn_opt.go

package info (click to toggle)
golang-github-sourcegraph-jsonrpc2 0.2.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 188 kB
  • sloc: makefile: 2
file content (111 lines) | stat: -rw-r--r-- 3,029 bytes parent folder | download
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
package jsonrpc2

import (
	"encoding/json"
	"sync"
)

// Logger interface implements one method - Printf.
// You can use the stdlib logger *log.Logger
type Logger interface {
	Printf(format string, v ...interface{})
}

// ConnOpt is the type of function that can be passed to NewConn to
// customize the Conn before it is created.
type ConnOpt func(*Conn)

// OnRecv causes all requests received on conn to invoke f(req, nil)
// and all responses to invoke f(req, resp),
func OnRecv(f func(*Request, *Response)) ConnOpt {
	return func(c *Conn) { c.onRecv = append(c.onRecv, f) }
}

// OnSend causes all requests sent on conn to invoke f(req, nil) and
// all responses to invoke f(nil, resp),
func OnSend(f func(*Request, *Response)) ConnOpt {
	return func(c *Conn) { c.onSend = append(c.onSend, f) }
}

// LogMessages causes all messages sent and received on conn to be
// logged using the provided logger.
func LogMessages(logger Logger) ConnOpt {
	return func(c *Conn) {
		// Remember reqs we have received so we can helpfully show the
		// request method in OnSend for responses.
		var (
			mu         sync.Mutex
			reqMethods = map[ID]string{}
		)

		// Set custom logger from provided input
		c.logger = logger

		OnRecv(func(req *Request, resp *Response) {
			switch {
			case req != nil:
				mu.Lock()
				reqMethods[req.ID] = req.Method
				mu.Unlock()

				params, _ := json.Marshal(req.Params)
				if req.Notif {
					logger.Printf("jsonrpc2: --> notif: %s: %s\n", req.Method, params)
				} else {
					logger.Printf("jsonrpc2: --> request #%s: %s: %s\n", req.ID, req.Method, params)
				}

			case resp != nil:
				var method string
				if req != nil {
					method = req.Method
				} else {
					method = "(no matching request)"
				}
				switch {
				case resp.Result != nil:
					result, _ := json.Marshal(resp.Result)
					logger.Printf("jsonrpc2: --> result #%s: %s: %s\n", resp.ID, method, result)
				case resp.Error != nil:
					err, _ := json.Marshal(resp.Error)
					logger.Printf("jsonrpc2: --> error #%s: %s: %s\n", resp.ID, method, err)
				}
			}
		})(c)
		OnSend(func(req *Request, resp *Response) {
			switch {
			case req != nil:
				params, _ := json.Marshal(req.Params)
				if req.Notif {
					logger.Printf("jsonrpc2: <-- notif: %s: %s\n", req.Method, params)
				} else {
					logger.Printf("jsonrpc2: <-- request #%s: %s: %s\n", req.ID, req.Method, params)
				}

			case resp != nil:
				mu.Lock()
				method := reqMethods[resp.ID]
				delete(reqMethods, resp.ID)
				mu.Unlock()
				if method == "" {
					method = "(no previous request)"
				}

				if resp.Result != nil {
					result, _ := json.Marshal(resp.Result)
					logger.Printf("jsonrpc2: <-- result #%s: %s: %s\n", resp.ID, method, result)
				} else {
					err, _ := json.Marshal(resp.Error)
					logger.Printf("jsonrpc2: <-- error #%s: %s: %s\n", resp.ID, method, err)
				}
			}
		})(c)
	}
}

// SetLogger sets the logger for the connection.
func SetLogger(logger Logger) ConnOpt {
	return func(c *Conn) {
		c.logger = logger
	}
}