chat server or client.

chat.go

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
)

type Client struct {
	token        string
	localWeather string
	conn         net.Conn
}

func (c *Client) Echo() {
	for {
		if c.token != "" {
			break
		}
		nameRd := bufio.NewReader(os.Stdin)
		fmt.Print("please named yourself: ")
		username, _ := nameRd.ReadString('\n')
		username = strings.TrimSpace(username)
		if username == "" {
			continue
		}
		c.token = username
		c.conn.Write([]byte(username))
	}
	log.Println("start chat...")
	for {
		rd := bufio.NewReader(os.Stdin)
		fmt.Printf(">> ")
		msg, _ := rd.ReadString('\n')
		msg = strings.TrimSpace(msg)
		if msg == "" {
			continue
		} else if msg == "?weather" {
			log.Println(c.Weather())
			continue
		}
		_, err := c.conn.Write([]byte(msg))
		if err != nil {
			log.Fatalln(err)
		}
		if msg == "?q" {
			c.conn.Close()
			fmt.Println("ByeBye..")
			os.Exit(0)
		}
	}
}

func (c *Client) Weather() string {
	if strings.TrimSpace(c.localWeather) != "" {
		return c.localWeather
	}
	c.localWeather = "getWeatherFromHTTPAPI"
	return c.localWeather
}

func (c *Client) Start(raddr string) {
	if c.conn == nil {
		rAddr, err := net.ResolveTCPAddr("tcp4", raddr)
		if err != nil {
			log.Fatalln(err)
		}
		c.conn, err = net.DialTCP("tcp", nil, rAddr)
		if err != nil {
			log.Fatalln(err)
		}
	}
	log.Println("Connecting to chat server...", raddr)
	go c.Echo()
	buf := make([]byte, 1024)
	for {
		n, err := c.conn.Read(buf)
		if err != nil {
			log.Fatalln(err)
		}
		log.Println(string(buf[:n]))
	}
}

type Server struct {
	listenAddr string
	clients    map[string]*Client
	joins      chan *M
	sends      chan *M
	lefts      chan *M
}

type M struct {
	c       *Client
	act     string
	message string
}

func (s *Server) broadcast(m *M) {
	for t, c := range s.clients {
		_, err := c.conn.Write([]byte(m.act + ": " + m.message))
		if err != nil {
			log.Println(err)
			delete(s.clients, t)
		}
	}
}

func (s *Server) Handler(c *Client) {
	raddr := c.conn.RemoteAddr().String()
	buft := make([]byte, 128)
	n, _ := c.conn.Read(buft)
	c.token = string(buft[:n])
	s.joins <- &M{c, "join", fmt.Sprintf("connection from %s as %s", raddr, c.token)}
	s.clients[c.token] = c
	buf := make([]byte, 1024)
	for {
		n, err := c.conn.Read(buf)
		if err != nil {
			log.Fatalln(err)
		}
		message := string(buf[:n])
		if message == "?q" {
			defer c.conn.Close()
			s.lefts <- &M{c, "left", fmt.Sprintf("%s disconnect", c.token)}
			delete(s.clients, c.token)
			return
		} else if message == "?users" {
			c.conn.Write([]byte(s.CurUsers()))
			continue
		}
		s.sends <- &M{c, "", fmt.Sprintf("%s says %s", c.token, message)}
	}
}

func (s *Server) CurUsers() string {
	buf := bytes.Buffer{}
	buf.WriteString("current users: ")
	for t := range s.clients {
		buf.WriteString(t + ", ")
	}
	return buf.String()
}

func (s *Server) BroadCast() {
	for {
		select {
		case join := <-s.joins:
			s.broadcast(join)
		case send := <-s.sends:
			s.broadcast(send)
		case left := <-s.lefts:
			s.broadcast(left)
		}
	}
}

func (s *Server) Start() {
	laddr, err := net.ResolveTCPAddr("tcp4", s.listenAddr)
	if err != nil {
		log.Fatalln(err)
	}
	l, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		log.Fatalln(err)
	}
	go s.BroadCast()
	log.Println("Listening on ", s.listenAddr)
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println(err)
			continue
		}
		c := &Client{conn: conn}
		go s.Handler(c)
	}
}

func main() {
	flag.Usage = func() {
		log.Fatalln("%s -role server/client -addr :8080", os.Args[0])
	}
	role := flag.String("role", "server", "play role, server or client")
	addr := flag.String("addr", ":8080", "Listening addr or remote addr")
	flag.Parse()
	if *role == "server" {
		s := &Server{
			listenAddr: *addr,
			clients:    make(map[string]*Client),
			joins:      make(chan *M, 1<<10),
			sends:      make(chan *M, 1<<10),
			lefts:      make(chan *M, 1<<10),
		}
		s.Start()
	} else if *role == "client" {
		c := &Client{}
		c.Start(*addr)
	} else {
		flag.Usage()
	}
}