sh3lly/shells/tcp.go

package shells

import (
	"bufio"
	"context"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"git.tcp.direct/kayos/sh3lly/chat"
	"git.tcp.direct/kayos/sh3lly/chat/message"
)

var srv *Server

// Server is an instance of our concurrent TCP server including a map of active retards
type Server struct {
	Map map[string]*Client
	mu  *sync.RWMutex
}

// Client represents a known retard
type Client struct {
	message.Identifier
	id   string
	Conn net.Conn

	chatUser   *message.User
	chatMember *chat.Member

	mu        *sync.RWMutex
	IP        string
	Group     string
	new       bool
	connected bool
	ctx       context.Context
	Cancel    context.CancelFunc
	read      *bufio.Reader
}

func closeConn(c *Client) {
	if err := c.Conn.Close(); err != nil {
		println(err.Error())
	}
	log.Warn().Msg("closed: " + c.Conn.RemoteAddr().String())
}

func (c *Client) recv(recvChan chan string) {
	for {
		select {
		case <-c.ctx.Done():
			c.connected = false
			go c.Conn.Close()
			go c.chatMember.User.Close()
			return
		default:
			if !c.connected {
				c.Cancel()
				return
			}
			in, err := c.read.ReadString('\n')
			if err != nil {
				log.Debug().Err(err).
					Str("caller", c.IP).
					Msg("failed to receive")
				c.connected = false
				continue
			}
			c.ctx, c.Cancel = context.WithDeadline(c.ctx, time.Now().Add(5*time.Minute))
			c.read.Reset(c.Conn)
			recvChan <- strings.ToLower(strings.TrimRight(in, "\n"))
		}
	}
}

func (c *Client) send(data string) {
	select {
	case <-c.ctx.Done():
		return
	default:
		if _, err := c.Conn.Write([]byte(data + "\n")); err != nil {
			c.connected = false
		}
	}
}

func (c *Client) newChatBot() (u *message.User, m *chat.Member) {
	var err error
	u = message.NewUser(c)
	u.SetConfig(message.UserConfig{Echo: false, Quiet: true, Zombie: true, ZombieConn: c.Conn, ZombieGrp: ""})

	var count = 0

	c.ctx, c.Cancel = context.WithDeadline(context.Background(), time.Now().Add(5*time.Minute))

	m, err = chat.MainRoom.Join(u)

	for err != nil {
		select {
		case <-c.ctx.Done():
			go m.Close()
			go u.Close()
			return
		default:
			if count > 10 {
				log.Error().Str("caller", c.Conn.RemoteAddr().String()).
					Err(err).Msg("giving up")
				go c.Cancel()
				return nil, nil
			}
			count++
			log.Error().Str("caller", c.Conn.RemoteAddr().String()).
				Err(err).Msg("failed to join, trying other ID")
			c.id = keygen()
			u.SetID(c.id)

			m, err = chat.MainRoom.Join(u)
		}
	}

	return
}

func (c *Client) setGroup(grp string) {
	grp = strings.TrimSpace(grp)
	c.mu.Lock()
	c.Group = grp
	c.chatUser.SetGroup(grp)
	chat.MainRoom.Send(message.NewPublicMsg("[internal] group id is now: "+c.Group, c.chatUser))
	c.mu.Unlock()
}

func (c *Client) handleIncoming() {
	groupprefix := "![NONE]!"

	setgrpprefix := fmt.Sprintf("!setgroup %s", c.id)
	setgrpprefixall := "!setgroup all"

	recvChan := make(chan string, 16)
	go c.recv(recvChan)

	defer func() {
		chat.MainRoom.Send(message.NewPublicMsg("%s has disconnected/timed out", c.chatUser))
		err := chat.MainRoom.Members.Remove(c.chatUser.Name())
		if err != nil {
			log.Warn().Str("caller", c.chatUser.Name()).Msg(err.Error())
		}
	}()

	for {
		prefix := fmt.Sprintf("!%s", c.id)

		if c.Group != "" {
			groupprefix = fmt.Sprintf("!%s", c.Group)
		}
		select {
		case <-c.ctx.Done():
			go c.chatUser.Close()
			go c.chatMember.Close()
			go c.Conn.Close()
			return
		case in := <-c.chatUser.ZombiePipe:
			in = strings.TrimSpace(in)
			switch {
			case strings.HasPrefix(in, prefix):
				c.send(strings.TrimPrefix(in, prefix))
				continue
			case strings.HasPrefix(in, "!all"):
				c.send(strings.TrimPrefix(in, "!all"))
				continue
			case strings.HasPrefix(in, groupprefix):
				c.send(strings.TrimPrefix(in, groupprefix))
				continue
			case strings.HasPrefix(in, setgrpprefixall):
				c.setGroup(strings.TrimPrefix(in, setgrpprefixall))
				continue
			case strings.HasPrefix(in, setgrpprefix):
				c.setGroup(strings.TrimPrefix(in, setgrpprefix))
				continue
			default:
				//
			}
		case in := <-recvChan:
			chat.MainRoom.Send(message.NewPublicMsg(in, c.chatUser))
		default:
			if !c.connected {
				return
			}
			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (s *Server) handleTCP(c *Client) {

	c.read = bufio.NewReader(c.Conn)
	if err := c.Conn.(*net.TCPConn).SetLinger(1); err != nil {
		fmt.Println("error while setting setlinger:", err.Error())
	}

	defer func() {
		content := fmt.Sprintf("%s has disconnected", c.Conn.RemoteAddr().String())
		chat.ExternalAnnouncements <- message.NewAnnounceMsg(content)
		go chat.MainRoom.Leave(c.chatUser)
		closeConn(c)
	}()

	c.chatUser, c.chatMember = c.newChatBot()

	go c.handleIncoming()

	for {
		if !c.connected {
			log.Debug().Str("caller", c.chatUser.ID()).Msg("disconnected")
			_ = chat.MainRoom.Leave(c.chatUser)
			c.chatUser.Close()
			return
		}
		c.chatUser.Consume()
		time.Sleep(100 * time.Millisecond)
	}
}