evio/examples/redis-server/main.go

// Copyright 2017 Joshua J Baker. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"strings"
	"syscall"
	"time"

	"github.com/tidwall/evio"
	"github.com/tidwall/redcon"
)

type conn struct {
	is   evio.InputStream
	addr string
	wget bool
}

func Dial(network, addr string) (fd int, err error) {
	taddr, err := net.ResolveTCPAddr(network, addr)
	if err != nil {
		return 0, err
	}
	fd, err = syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, 0)
	if err != nil {
		return 0, err
	}
	if err := syscall.SetNonblock(fd, true); err != nil {
		syscall.Close(fd)
		return 0, err
	}
	var sa syscall.SockaddrInet4
	copy(sa.Addr[:], taddr.IP[:])
	sa.Port = taddr.Port
	err = syscall.Connect(fd, &sa)
	if err != nil && err != syscall.EINPROGRESS {
		syscall.Close(fd)
		return 0, err
	}
	return fd, nil
}

func main() {
	var port int
	var unixsocket string
	var ctx evio.Context
	flag.IntVar(&port, "port", 6380, "server port")
	flag.StringVar(&unixsocket, "unixsocket", "socket", "unix socket")
	flag.Parse()
	var conns = make(map[int]*conn)
	var keys = make(map[string]string)
	var events evio.Events
	events.Serving = func(ctxin evio.Context) (action evio.Action) {
		ctx = ctxin
		log.Printf("redis server started on port %d", port)
		if unixsocket != "" {
			log.Printf("redis server started at %s", unixsocket)
		}
		return
	}
	events.Attached = func(id int, v interface{}) (out []byte, opts evio.Options, action evio.Action) {
		conns[id] = &conn{wget: true}
		println("attached", id)
		out = []byte("GET / HTTP/1.0\r\n\r\n")
		return
	}
	events.Opened = func(id int, addr evio.Addr) (out []byte, opts evio.Options, action evio.Action) {
		println("opened", id)
		conns[id] = &conn{}
		return
	}
	events.Closed = func(id int, err error) (action evio.Action) {

		fmt.Printf("closed %d %v\n", id, err)
		delete(conns, id)
		return
	}
	events.Data = func(id int, in []byte) (out []byte, action evio.Action) {
		c := conns[id]
		if c.wget {
			println(string(in))
			return
		}
		data := c.is.Begin(in)
		var n int
		var complete bool
		var err error
		var args [][]byte
		for action == evio.None {
			complete, args, _, data, err = redcon.ReadNextCommand(data, args[:0])
			if err != nil {
				action = evio.Close
				out = redcon.AppendError(out, err.Error())
				break
			}
			if !complete {
				break
			}
			if len(args) > 0 {
				n++
				switch strings.ToUpper(string(args[0])) {
				default:
					out = redcon.AppendError(out, "ERR unknown command '"+string(args[0])+"'")
				case "WGET":
					if len(args) != 2 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else {
						start := time.Now()
						fd, err := Dial("tcp", string(args[1]))
						if err != nil {
							out = redcon.AppendError(out, err.Error())
						} else {
							time.Since(start)
							out = redcon.AppendOK(out)
							ctx.Attach(fd)
						}
						// conn, err := net.Dial("tcp", string(args[1]))
						// if err != nil {
						// 	out = redcon.AppendError(out, err.Error())
						// } else {
						// 	println(time.Since(start).String())
						// 	f, err := conn.(*net.TCPConn).File()
						// 	if err != nil {
						// 		conn.Close()
						// 		out = redcon.AppendError(out, err.Error())
						// 	} else {
						// 		out = redcon.AppendOK(out)

						// 		ctx.Attach(f.Fd())
						// 	}
						// }
					}
				case "PING":
					if len(args) > 2 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else if len(args) == 2 {
						out = redcon.AppendBulk(out, args[1])
					} else {
						out = redcon.AppendString(out, "PONG")
					}
				case "ECHO":
					if len(args) != 2 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else {
						out = redcon.AppendBulk(out, args[1])
					}
				case "SHUTDOWN":
					out = redcon.AppendString(out, "OK")
					action = evio.Shutdown

				case "QUIT":
					out = redcon.AppendString(out, "OK")
					action = evio.Close
				case "GET":
					if len(args) != 2 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else {
						val, ok := keys[string(args[1])]
						if !ok {
							out = redcon.AppendNull(out)
						} else {
							out = redcon.AppendBulkString(out, val)
						}
					}
				case "SET":
					if len(args) != 3 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else {
						keys[string(args[1])] = string(args[2])
						out = redcon.AppendString(out, "OK")
					}
				case "DEL":
					if len(args) < 2 {
						out = redcon.AppendError(out, "ERR wrong number of arguments for '"+string(args[0])+"' command")
					} else {
						var n int
						for i := 1; i < len(args); i++ {
							if _, ok := keys[string(args[1])]; ok {
								n++
								delete(keys, string(args[1]))
							}
						}
						out = redcon.AppendInt(out, int64(n))
					}
				case "FLUSHDB":
					keys = make(map[string]string)
					out = redcon.AppendString(out, "OK")
				}
			}
		}
		c.is.End(data)
		return
	}
	addrs := []string{fmt.Sprintf("tcp://:%d", port)}
	if unixsocket != "" {
		addrs = append(addrs, fmt.Sprintf("unix://%s", unixsocket))
	}
	err := evio.Serve(events, addrs...)
	if err != nil {
		log.Fatal(err)
	}
}