writer.go raw

   1  package io
   2  
   3  import (
   4  	"net"
   5  	"time"
   6  
   7  	"git.smesh.lol/torly/pkg/wire"
   8  	"git.smesh.lol/torly/pkg/ws"
   9  )
  10  
  11  const writeTimeout = 10 * time.Second
  12  
  13  // Writer is one writer worker domain. It owns the write side and the close of
  14  // the connection currently assigned to it. It never reads application data.
  15  func Writer(_ WriterData, ch WriterChans) {
  16  	var sock *wire.Socket
  17  	wsMode := false
  18  	closeConn := func() {
  19  		if sock != nil {
  20  			sock.Conn.Close()
  21  			sock = nil
  22  		}
  23  		wsMode = false
  24  	}
  25  	defer closeConn()
  26  
  27  	for {
  28  		select {
  29  		case <-ch.Stop.Done():
  30  			return
  31  		case ctl := <-ch.Ctl.Recv():
  32  			switch ctl.Kind {
  33  			case wire.CtlBind:
  34  				closeConn()
  35  				sock = ctl.Socket
  36  				wsMode = false
  37  				sendAck(ch, wire.CtlBind)
  38  			case wire.CtlUpgrade:
  39  				if sock == nil {
  40  					continue
  41  				}
  42  				if err := writeAll(sock.Conn, ctl.Payload); err != nil {
  43  					closeConn()
  44  					continue
  45  				}
  46  				wsMode = true
  47  				sendAck(ch, wire.CtlUpgrade)
  48  			case wire.CtlWrite:
  49  				if sock != nil {
  50  					if err := writeAll(sock.Conn, ctl.Payload); err != nil {
  51  						closeConn()
  52  					}
  53  				}
  54  			case wire.CtlPong:
  55  				if sock != nil && wsMode {
  56  					if err := writeAll(sock.Conn, ws.AppendFrame(nil, ws.OpPong, ctl.Payload)); err != nil {
  57  						closeConn()
  58  					}
  59  				}
  60  			case wire.CtlRespond:
  61  				if sock != nil {
  62  					_ = writeAll(sock.Conn, ctl.Payload)
  63  				}
  64  				closeConn()
  65  			case wire.CtlClose:
  66  				if sock != nil && wsMode {
  67  					frame := ctl.Payload
  68  					if len(frame) == 0 {
  69  						frame = ws.AppendClose(nil, 1000, "")
  70  					}
  71  					_ = writeAll(sock.Conn, frame)
  72  				}
  73  				closeConn()
  74  			}
  75  		case out := <-ch.Out.Recv():
  76  			if sock == nil {
  77  				continue
  78  			}
  79  			switch out.Kind {
  80  			case wire.OutText:
  81  				if wsMode {
  82  					if err := writeAll(sock.Conn, ws.AppendFrame(nil, ws.OpText, out.Text)); err != nil {
  83  						closeConn()
  84  					}
  85  				}
  86  			case wire.OutHTTP:
  87  				_ = writeAll(sock.Conn, SerializeHTTP(out.HTTP))
  88  				closeConn()
  89  			case wire.OutClose:
  90  				if wsMode {
  91  					frame := ws.AppendClose(nil, out.Code, out.Reason)
  92  					_ = writeAll(sock.Conn, frame)
  93  				}
  94  				closeConn()
  95  			case wire.OutPing:
  96  				if wsMode {
  97  					if err := writeAll(sock.Conn, ws.AppendFrame(nil, ws.OpPing, nil)); err != nil {
  98  						closeConn()
  99  					}
 100  				}
 101  			}
 102  		}
 103  	}
 104  }
 105  
 106  func sendAck(ch WriterChans, kind uint8) {
 107  	select {
 108  	case ch.Ack.Tx <- kind:
 109  	case <-ch.Stop.Done():
 110  	}
 111  }
 112  
 113  func writeAll(c net.Conn, p []byte) error {
 114  	_ = c.SetWriteDeadline(time.Now().Add(writeTimeout))
 115  	for len(p) > 0 {
 116  		n, err := c.Write(p)
 117  		if err != nil {
 118  			return err
 119  		}
 120  		p = p[n:]
 121  	}
 122  	return nil
 123  }
 124