[210] | 1 | package soju
|
---|
| 2 |
|
---|
| 3 | import (
|
---|
| 4 | "fmt"
|
---|
| 5 | "net"
|
---|
| 6 | "time"
|
---|
| 7 |
|
---|
| 8 | "gopkg.in/irc.v3"
|
---|
| 9 | )
|
---|
| 10 |
|
---|
| 11 | func setKeepAlive(c net.Conn) error {
|
---|
| 12 | tcpConn, ok := c.(*net.TCPConn)
|
---|
| 13 | if !ok {
|
---|
| 14 | return fmt.Errorf("cannot enable keep-alive on a non-TCP connection")
|
---|
| 15 | }
|
---|
| 16 | if err := tcpConn.SetKeepAlive(true); err != nil {
|
---|
| 17 | return err
|
---|
| 18 | }
|
---|
| 19 | return tcpConn.SetKeepAlivePeriod(keepAlivePeriod)
|
---|
| 20 | }
|
---|
| 21 |
|
---|
| 22 | type conn struct {
|
---|
| 23 | net net.Conn
|
---|
| 24 | irc *irc.Conn
|
---|
| 25 | srv *Server
|
---|
| 26 | logger Logger
|
---|
| 27 | outgoing chan<- *irc.Message
|
---|
| 28 | closed chan struct{}
|
---|
| 29 | }
|
---|
| 30 |
|
---|
| 31 | func newConn(srv *Server, netConn net.Conn, logger Logger) *conn {
|
---|
| 32 | setKeepAlive(netConn)
|
---|
| 33 |
|
---|
| 34 | outgoing := make(chan *irc.Message, 64)
|
---|
| 35 | c := &conn{
|
---|
| 36 | net: netConn,
|
---|
| 37 | irc: irc.NewConn(netConn),
|
---|
| 38 | srv: srv,
|
---|
| 39 | outgoing: outgoing,
|
---|
| 40 | logger: logger,
|
---|
| 41 | closed: make(chan struct{}),
|
---|
| 42 | }
|
---|
| 43 |
|
---|
| 44 | go func() {
|
---|
| 45 | for msg := range outgoing {
|
---|
| 46 | if c.srv.Debug {
|
---|
| 47 | c.logger.Printf("sent: %v", msg)
|
---|
| 48 | }
|
---|
| 49 | c.net.SetWriteDeadline(time.Now().Add(writeTimeout))
|
---|
| 50 | if err := c.irc.WriteMessage(msg); err != nil {
|
---|
| 51 | c.logger.Printf("failed to write message: %v", err)
|
---|
| 52 | break
|
---|
| 53 | }
|
---|
| 54 | }
|
---|
| 55 | if err := c.net.Close(); err != nil {
|
---|
| 56 | c.logger.Printf("failed to close connection: %v", err)
|
---|
| 57 | } else {
|
---|
| 58 | c.logger.Printf("connection closed")
|
---|
| 59 | }
|
---|
| 60 | // Drain the outgoing channel to prevent SendMessage from blocking
|
---|
| 61 | for range outgoing {
|
---|
| 62 | // This space is intentionally left blank
|
---|
| 63 | }
|
---|
| 64 | }()
|
---|
| 65 |
|
---|
| 66 | c.logger.Printf("new connection")
|
---|
| 67 | return c
|
---|
| 68 | }
|
---|
| 69 |
|
---|
| 70 | func (c *conn) isClosed() bool {
|
---|
| 71 | select {
|
---|
| 72 | case <-c.closed:
|
---|
| 73 | return true
|
---|
| 74 | default:
|
---|
| 75 | return false
|
---|
| 76 | }
|
---|
| 77 | }
|
---|
| 78 |
|
---|
| 79 | // Close closes the connection. It is safe to call from any goroutine.
|
---|
| 80 | func (c *conn) Close() error {
|
---|
| 81 | if c.isClosed() {
|
---|
| 82 | return fmt.Errorf("connection already closed")
|
---|
| 83 | }
|
---|
| 84 | close(c.closed)
|
---|
| 85 | close(c.outgoing)
|
---|
| 86 | return nil
|
---|
| 87 | }
|
---|
| 88 |
|
---|
| 89 | func (c *conn) ReadMessage() (*irc.Message, error) {
|
---|
| 90 | msg, err := c.irc.ReadMessage()
|
---|
| 91 | if err != nil {
|
---|
| 92 | return nil, err
|
---|
| 93 | }
|
---|
| 94 |
|
---|
| 95 | if c.srv.Debug {
|
---|
| 96 | c.logger.Printf("received: %v", msg)
|
---|
| 97 | }
|
---|
| 98 |
|
---|
| 99 | return msg, nil
|
---|
| 100 | }
|
---|
| 101 |
|
---|
| 102 | // SendMessage queues a new outgoing message. It is safe to call from any
|
---|
| 103 | // goroutine.
|
---|
| 104 | func (c *conn) SendMessage(msg *irc.Message) {
|
---|
| 105 | if c.isClosed() {
|
---|
| 106 | return
|
---|
| 107 | }
|
---|
| 108 | c.outgoing <- msg
|
---|
| 109 | }
|
---|