github.com/fozzysec/SiaPrime@v0.0.0-20190612043147-66c8e8d11fe3/modules/stratumminer/tcpclient.go (about) 1 // Package stratumminer implements the basic stratum protocol. 2 // This is normal jsonrpc but the go standard library is insufficient since we need features like notifications. 3 package stratumminer 4 5 import ( 6 "bufio" 7 "encoding/json" 8 "errors" 9 "fmt" 10 "log" 11 "net" 12 "sync" 13 "time" 14 15 //"SiaPrime/build" 16 "SiaPrime/types" 17 18 "gitlab.com/NebulousLabs/threadgroup" 19 ) 20 21 //ErrorCallback is the type of function that be registered to be notified of errors requiring a client 22 // to be dropped and a new one to be created 23 type ErrorCallback func(err error) 24 25 //NotificationHandler is the signature for a function that handles notifications 26 type NotificationHandler func(args []interface{}) 27 28 // TcpClient maintains a connection to the stratum server and (de)serializes requests/reponses/notifications 29 type TcpClient struct { 30 socket net.Conn 31 mu sync.Mutex // protects connected state 32 connected bool 33 34 seqmutex sync.Mutex // protects following 35 seq uint64 36 37 callsMutex sync.Mutex // protects following 38 pendingCalls map[uint64]chan interface{} 39 40 ErrorCallback ErrorCallback 41 notificationHandlers map[string]NotificationHandler 42 43 tg threadgroup.ThreadGroup 44 } 45 46 // Dial connects to a stratum+tcp at the specified network address. 47 // This function is not threadsafe 48 // If an error occurs, it is both returned here and through the ErrorCallback of the TcpClient 49 func (c *TcpClient) Dial(host string) (err error) { 50 c.mu.Lock() 51 c.connected = false 52 c.mu.Unlock() 53 54 select { 55 case <-c.tg.StopChan(): 56 return 57 default: 58 } 59 60 log.Println("TcpClient Dialing") 61 c.mu.Lock() 62 c.connected = false 63 c.socket, err = net.Dial("tcp", host) 64 c.mu.Unlock() 65 select { 66 case <-c.tg.StopChan(): 67 return 68 default: 69 } 70 if err != nil { 71 //log.Println(err) 72 c.dispatchError(err) 73 return err 74 } 75 c.tg.OnStop(func() error { 76 log.Println("TCPClient: Closing c.socket") 77 c.cancelAllRequests() 78 c.socket.Close() 79 return nil 80 }) 81 c.connected = true 82 //log.Println("TcpClient Done Dialing") 83 go c.Listen() 84 return 85 } 86 87 // Close releases the tcp connection 88 func (c *TcpClient) Close() { 89 log.Println("TcpClient Close() called") 90 if err := c.tg.Stop(); err != nil { 91 panic(err.Error()) 92 } 93 log.Println("Closing TcpClient socket") 94 c.mu.Lock() 95 defer c.mu.Unlock() 96 c.connected = false 97 //log.Println("Done closing TcpClient") 98 } 99 100 // Connected returns whether or not the tcp client has an open tcp socket 101 func (c *TcpClient) Connected() bool { 102 c.mu.Lock() 103 defer c.mu.Unlock() 104 log.Printf("Checking if connected: %t\n", c.connected) 105 return c.connected 106 } 107 108 // SetNotificationHandler registers a function to handle notification for a specific method. 109 // This function is not threadsafe and all notificationhandlers should be set prior to calling the Dial function 110 func (c *TcpClient) SetNotificationHandler(method string, handler NotificationHandler) { 111 if c.notificationHandlers == nil { 112 c.notificationHandlers = make(map[string]NotificationHandler) 113 } 114 c.notificationHandlers[method] = handler 115 } 116 117 func (c *TcpClient) dispatchNotification(n types.StratumNotification) { 118 if c.notificationHandlers == nil { 119 return 120 } 121 if notificationHandler, exists := c.notificationHandlers[n.Method]; exists { 122 notificationHandler(n.Params) 123 } 124 } 125 126 func (c *TcpClient) dispatch(r types.StratumResponse) { 127 if r.ID == 0 { 128 log.Println("Dispatching notification") 129 log.Println(r.StratumNotification) 130 c.dispatchNotification(r.StratumNotification) 131 return 132 } 133 c.callsMutex.Lock() 134 defer c.callsMutex.Unlock() 135 cb, found := c.pendingCalls[r.ID] 136 var result interface{} 137 //log.Printf("dispatch response: %s\n", r.Result) 138 //log.Printf("dispatch error: %s\n", r.Error) 139 if r.Error != nil { 140 message := "" 141 if len(r.Error) >= 2 { 142 message, _ = r.Error[1].(string) 143 } 144 result = errors.New(message) 145 /* 146 var message []byte 147 r.Error.UnmarshalJSON(message) 148 result = errors.New(string(message[:])) 149 */ 150 } else { 151 result = r.Result 152 } 153 if found { 154 cb <- result 155 } 156 } 157 158 func (c *TcpClient) dispatchError(err error) { 159 log.Println("dispatching error") 160 select { 161 // don't dispatch any errors if we've been shutdown! 162 case <-c.tg.StopChan(): 163 log.Println("stop called, not dispatching error") 164 return 165 default: 166 } 167 if c.ErrorCallback != nil { 168 c.ErrorCallback(err) 169 } 170 } 171 172 //Listen reads data from the open connection, deserializes it and dispatches the reponses and notifications 173 // This is a blocking function and will continue to listen until an error occurs (io or deserialization) 174 func (c *TcpClient) Listen() { 175 /* 176 if err := c.tg.Add(); err != nil { 177 build.Critical(err) 178 } 179 defer c.tg.Done() 180 */ 181 reader := bufio.NewReader(c.socket) 182 for { 183 rawmessage, err := reader.ReadString('\n') 184 // bail out if we've called stop 185 select { 186 case <-c.tg.StopChan(): 187 log.Println("TCPCLIENT StopChan called, done Listen()ing") 188 return 189 default: 190 } 191 if err != nil { 192 log.Printf("TCPCLIENT ERR: %s\n", err) 193 c.dispatchError(err) 194 return 195 } 196 r := types.StratumResponse{} 197 err = json.Unmarshal([]byte(rawmessage), &r) 198 if err != nil { 199 c.dispatchError(err) 200 return 201 } 202 c.dispatch(r) 203 } 204 } 205 206 func (c *TcpClient) registerRequest(requestID uint64) (cb chan interface{}) { 207 c.callsMutex.Lock() 208 defer c.callsMutex.Unlock() 209 if c.pendingCalls == nil { 210 c.pendingCalls = make(map[uint64]chan interface{}) 211 } 212 cb = make(chan interface{}) 213 c.pendingCalls[requestID] = cb 214 return 215 } 216 217 func (c *TcpClient) cancelRequest(requestID uint64) { 218 c.callsMutex.Lock() 219 defer c.callsMutex.Unlock() 220 cb, found := c.pendingCalls[requestID] 221 if found { 222 close(cb) 223 delete(c.pendingCalls, requestID) 224 } 225 } 226 227 func (c *TcpClient) cancelAllRequests() { 228 c.callsMutex.Lock() 229 defer c.callsMutex.Unlock() 230 for requestID, cb := range c.pendingCalls { 231 close(cb) 232 delete(c.pendingCalls, requestID) 233 } 234 } 235 236 //Call invokes the named function, waits for it to complete, and returns its error status. 237 func (c *TcpClient) Call(serviceMethod string, args []string) (reply interface{}, err error) { 238 //jsonargs, _ := json.Marshal(args) 239 //rawargs := json.RawMessage(jsonargs) 240 params := make([]interface{}, len(args)) 241 for i, v := range args { 242 params[i] = v 243 } 244 r := types.StratumRequest{Method: serviceMethod, Params: params} 245 //r := types.StratumRequest{Method: serviceMethod, Params: rawargs} 246 247 c.seqmutex.Lock() 248 c.seq++ 249 r.ID = c.seq 250 c.seqmutex.Unlock() 251 252 rawmsg, err := json.Marshal(r) 253 if err != nil { 254 err = fmt.Errorf("json.Marshal failed: %v", err) 255 return 256 } 257 call := c.registerRequest(r.ID) 258 defer c.cancelRequest(r.ID) 259 260 rawmsg = append(rawmsg, []byte("\n")...) 261 c.mu.Lock() 262 if c.connected { 263 _, err = c.socket.Write(rawmsg) 264 } else { 265 err = fmt.Errorf("Can't write to socket, socket has been closed") 266 return nil, err 267 } 268 c.mu.Unlock() 269 if err != nil { 270 err = fmt.Errorf("socket.Write failed: %v", err) 271 return 272 } 273 //Make sure the request is cancelled if no response is given 274 go func() { 275 // cancel after 10 seconds 276 for timeElapsed := 0; timeElapsed < 10; timeElapsed += 1 { 277 // cancel the request if we've called stop 278 select { 279 case <-c.tg.StopChan(): 280 return 281 default: 282 time.Sleep(1 * time.Second) 283 } 284 } 285 c.cancelRequest(r.ID) 286 }() 287 reply = <-call 288 289 if reply == nil { 290 err = errors.New("Timeout") 291 return 292 } 293 err, _ = reply.(error) 294 return 295 }