<commit_before>package main

import (
	"crypto/sha1"
	"fmt"
	"io"
	"log"
	"net"
	"regexp"
	"strings"
)

var (
	nickRegexp    = regexp.MustCompile(`^[a-zA-Z\[\]_^{|}][a-zA-Z0-9\[\]_^{|}]*$`)
	channelRegexp = regexp.MustCompile(`^#[a-zA-Z0-9_\-]+$`)
)

func NewServer() *Server {
	return &Server{eventChan: make(chan Event),
		name:        "rosella",
		clientMap:   make(map[string]*Client),
		channelMap:  make(map[string]*Channel),
		operatorMap: make(map[string]string),
		motd:        "Welcome to IRC. Powered by Rosella."}
}

func (s *Server) Run() {
	for event := range s.eventChan {
		s.handleEvent(event)
	}
}

func (s *Server) HandleConnection(conn net.Conn) {
	client := &Client{server: s,
		connection: conn,
		outputChan: make(chan string),
		signalChan: make(chan signalCode, 3),
		channelMap: make(map[string]*Channel),
		connected:  true}

	go client.clientThread()
}

func (s *Server) handleEvent(e Event) {
	defer func(event Event) {
		err := recover()
		if err != nil {
			log.Printf("Recovered from error when handling event: %+v", event)
			log.Println(err)
		}
	}(e)

	switch e.event {
	case connected:
		//Client connected
		e.client.reply(rplMOTD, s.motd)
	case disconnected:
		//Client disconnected
	case command:
		//Client send a command
		fields := strings.Fields(e.input)
		if len(fields) < 1 {
			return
		}

		if strings.HasPrefix(fields[0], ":") {
			fields = fields[1:]
		}
		command := strings.ToUpper(fields[0])
		args := fields[1:]

		s.handleCommand(e.client, command, args)
	}
}

func (s *Server) handleCommand(client *Client, command string, args []string) {

	switch command {
	case "PING":
		client.reply(rplPong)
	case "INFO":
		client.reply(rplInfo, "Rosella IRCD github.com/eXeC64/Rosella")
	case "VERSION":
		client.reply(rplVersion, VERSION)
	case "NICK":
		if len(args) < 1 {
			client.reply(errNoNick)
			return
		}

		newNick := args[0]

		//Check newNick is of valid formatting (regex)
		if nickRegexp.MatchString(newNick) == false {
			client.reply(errInvalidNick, newNick)
			return
		}

		if _, exists := s.clientMap[strings.ToLower(newNick)]; exists {
			client.reply(errNickInUse, newNick)
			return
		}

		//Protect the server name from being used
		if strings.ToLower(newNick) == strings.ToLower(s.name) {
			client.reply(errNickInUse, newNick)
			return
		}

		client.setNick(newNick)

	case "USER":
		if client.nick == "" {
			client.reply(rplKill, "Your nickname is already being used", "")
			client.disconnect()
		} else {
			client.reply(rplWelcome)
			client.registered = true
		}

	case "JOIN":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		if args[0] == "0" {
			//Quit all channels
			for channel := range client.channelMap {
				client.partChannel(channel, "Disconnecting")
			}
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Join the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				client.joinChannel(channel)
			}
		}

	case "PART":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		reason := strings.Join(args[1:], " ")

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Part the channel if it's valid
			if channelRegexp.Match([]byte(channel)) {
				client.partChannel(channel, reason)
			}
		}

	case "PRIVMSG":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		message := strings.Join(args[1:], " ")

		channel, chanExists := s.channelMap[strings.ToLower(args[0])]
		client2, clientExists := s.clientMap[strings.ToLower(args[0])]

		if chanExists {
			if channel.mode.noExternal {
				if _, inChannel := channel.clientMap[strings.ToLower(client.nick)]; !inChannel {
					//Not in channel, not allowed to send
					client.reply(errCannotSend, args[0])
					return
				}
			}
			if channel.mode.moderated {
				clientMode := channel.modeMap[strings.ToLower(client.nick)]
				if !clientMode.operator && !clientMode.voice {
					//It's moderated and we're not +v or +o, do nothing
					client.reply(errCannotSend, args[0])
					return
				}
			}
			for _, c := range channel.clientMap {
				if c != client {
					c.reply(rplMsg, client.nick, args[0], message)
				}
			}
		} else if clientExists {
			client.reply(rplMsg, client.nick, client2.nick, message)
		} else {
			client.reply(errNoSuchNick, args[0])
		}

	case "QUIT":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		client.disconnect()

	case "TOPIC":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		channel, exists := s.channelMap[strings.ToLower(args[0])]
		if exists == false {
			client.reply(errNoSuchNick, args[0])
			return
		}

		if len(args) == 1 {
			client.reply(rplTopic, channel.name, channel.topic)
			return
		}

		clientMode := channel.modeMap[strings.ToLower(client.nick)]
		if channel.mode.topicLocked && !clientMode.operator {
			client.reply(errNoPriv)
			return
		}

		if args[1] == ":" {
			channel.topic = ""
			for _, client := range channel.clientMap {
				client.reply(rplNoTopic, channel.name)
			}
		} else {
			topic := strings.Join(args[1:], " ")
			topic = strings.TrimPrefix(topic, ":")
			channel.topic = topic

			for _, client := range channel.clientMap {
				client.reply(rplTopic, channel.name, channel.topic)
			}
		}

	case "LIST":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) == 0 {
			chanList := make([]string, 0, len(s.channelMap))

			for channelName, channel := range s.channelMap {
				if channel.mode.secret {
					if _, inChannel := channel.clientMap[strings.ToLower(client.nick)]; !inChannel {
						//Not in the channel, skip
						continue
					}
				}
				listItem := fmt.Sprintf("%s %d :%s", channelName, len(channel.clientMap), channel.topic)
				chanList = append(chanList, listItem)
			}

			client.reply(rplList, chanList...)

		} else {
			channels := strings.Split(args[0], ",")
			chanList := make([]string, 0, len(channels))

			for _, channelName := range channels {
				if channel, exists := s.channelMap[strings.ToLower(channelName)]; exists {
					listItem := fmt.Sprintf("%s %d :%s", channelName, len(channel.clientMap), channel.topic)
					chanList = append(chanList, listItem)
				}
			}

			client.reply(rplList, chanList...)
		}
	case "OPER":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		username := args[0]
		password := args[1]

		if hashedPassword, exists := s.operatorMap[username]; exists {
			h := sha1.New()
			io.WriteString(h, password)
			pass := fmt.Sprintf("%x", h.Sum(nil))
			if hashedPassword == pass {
				client.operator = true
				client.reply(rplOper)
				return
			}
		}
		client.reply(errPassword)

	case "KILL":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if client.operator == false {
			client.reply(errNoPriv)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		nick := args[0]

		reason := strings.Join(args[1:], " ")

		client, exists := s.clientMap[strings.ToLower(nick)]
		if !exists {
			client.reply(errNoSuchNick, nick)
			return
		}

		client.reply(rplKill, client.nick, reason)
		client.disconnect()

	case "KICK":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		channelKey := strings.ToLower(args[0])
		targetKey := strings.ToLower(args[1])

		channel, channelExists := s.channelMap[channelKey]
		if !channelExists {
			client.reply(errNoSuchNick, args[0])
			return
		}

		target, targetExists := channel.clientMap[targetKey]
		if !targetExists {
			client.reply(errNoSuchNick, args[1])
			return
		}

		clientMode := channel.modeMap[client.key]
		if !clientMode.operator && !client.operator {
			client.reply(errNoPriv)
			return
		}

		reason := strings.Join(args[2:], " ")

		//It worked
		for _, client := range channel.clientMap {
			client.reply(rplKick, client.nick, channel.name, target.nick, reason)
		}

		delete(channel.clientMap, targetKey)
		delete(channel.modeMap, targetKey)
		delete(target.channelMap, channelKey)

	case "MODE":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		channelKey := strings.ToLower(args[0])

		channel, channelExists := s.channelMap[channelKey]
		if !channelExists {
			client.reply(errNoSuchNick, args[0])
			return
		}
		mode := channel.mode

		if len(args) == 1 {
			//No more args, they just want the mode
			client.reply(rplChannelModeIs, args[0], mode.String(), "")
			return
		}

		if cm, ok := channel.modeMap[strings.ToLower(client.nick)]; !ok || !cm.operator {
			//Not a channel operator.

			//If they're not an irc operator either, they'll fail
			if !client.operator {
				client.reply(errNoPriv)
				return
			}
		}

		hasClient := false
		var oldClientMode, newClientMode *ClientMode
		var targetClient *Client
		if len(args) >= 3 {
			clientKey := strings.ToLower(args[2])
			oldClientMode, hasClient = channel.modeMap[clientKey]
			if hasClient {
				targetClient = channel.clientMap[clientKey]
				newClientMode = new(ClientMode)
				*newClientMode = *oldClientMode
			}
		}

		mod := strings.ToLower(args[1])
		if strings.HasPrefix(mod, "+") {
			for _, char := range mod {
				switch char {
				case 's':
					mode.secret = true
				case 't':
					mode.topicLocked = true
				case 'm':
					mode.moderated = true
				case 'n':
					mode.noExternal = true
				case 'o':
					if hasClient {
						newClientMode.operator = true
					}
				case 'v':
					if hasClient {
						newClientMode.voice = true
					}
				}
			}
		} else if strings.HasPrefix(mod, "-") {
			for _, char := range mod {
				switch char {
				case 's':
					mode.secret = false
				case 't':
					mode.topicLocked = false
				case 'm':
					mode.moderated = false
				case 'n':
					mode.noExternal = false
				case 'o':
					if hasClient {
						newClientMode.operator = false
					}
				case 'v':
					if hasClient {
						newClientMode.voice = false
					}
				}
			}
		}

		if hasClient {
			*oldClientMode = *newClientMode
		}
		channel.mode = mode

		for _, client := range channel.clientMap {
			if hasClient {
				client.reply(rplChannelModeIs, channel.name, args[1], targetClient.nick)
			} else {
				client.reply(rplChannelModeIs, channel.name, args[1], "")
			}
		}

	default:
		client.reply(errUnknownCommand, command)
	}
}
<commit_msg>Use MatchString instead of Match to save a []byte cast<commit_after>package main

import (
	"crypto/sha1"
	"fmt"
	"io"
	"log"
	"net"
	"regexp"
	"strings"
)

var (
	nickRegexp    = regexp.MustCompile(`^[a-zA-Z\[\]_^{|}][a-zA-Z0-9\[\]_^{|}]*$`)
	channelRegexp = regexp.MustCompile(`^#[a-zA-Z0-9_\-]+$`)
)

func NewServer() *Server {
	return &Server{eventChan: make(chan Event),
		name:        "rosella",
		clientMap:   make(map[string]*Client),
		channelMap:  make(map[string]*Channel),
		operatorMap: make(map[string]string),
		motd:        "Welcome to IRC. Powered by Rosella."}
}

func (s *Server) Run() {
	for event := range s.eventChan {
		s.handleEvent(event)
	}
}

func (s *Server) HandleConnection(conn net.Conn) {
	client := &Client{server: s,
		connection: conn,
		outputChan: make(chan string),
		signalChan: make(chan signalCode, 3),
		channelMap: make(map[string]*Channel),
		connected:  true}

	go client.clientThread()
}

func (s *Server) handleEvent(e Event) {
	defer func(event Event) {
		err := recover()
		if err != nil {
			log.Printf("Recovered from error when handling event: %+v", event)
			log.Println(err)
		}
	}(e)

	switch e.event {
	case connected:
		//Client connected
		e.client.reply(rplMOTD, s.motd)
	case disconnected:
		//Client disconnected
	case command:
		//Client send a command
		fields := strings.Fields(e.input)
		if len(fields) < 1 {
			return
		}

		if strings.HasPrefix(fields[0], ":") {
			fields = fields[1:]
		}
		command := strings.ToUpper(fields[0])
		args := fields[1:]

		s.handleCommand(e.client, command, args)
	}
}

func (s *Server) handleCommand(client *Client, command string, args []string) {

	switch command {
	case "PING":
		client.reply(rplPong)
	case "INFO":
		client.reply(rplInfo, "Rosella IRCD github.com/eXeC64/Rosella")
	case "VERSION":
		client.reply(rplVersion, VERSION)
	case "NICK":
		if len(args) < 1 {
			client.reply(errNoNick)
			return
		}

		newNick := args[0]

		//Check newNick is of valid formatting (regex)
		if nickRegexp.MatchString(newNick) == false {
			client.reply(errInvalidNick, newNick)
			return
		}

		if _, exists := s.clientMap[strings.ToLower(newNick)]; exists {
			client.reply(errNickInUse, newNick)
			return
		}

		//Protect the server name from being used
		if strings.ToLower(newNick) == strings.ToLower(s.name) {
			client.reply(errNickInUse, newNick)
			return
		}

		client.setNick(newNick)

	case "USER":
		if client.nick == "" {
			client.reply(rplKill, "Your nickname is already being used", "")
			client.disconnect()
		} else {
			client.reply(rplWelcome)
			client.registered = true
		}

	case "JOIN":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		if args[0] == "0" {
			//Quit all channels
			for channel := range client.channelMap {
				client.partChannel(channel, "Disconnecting")
			}
			return
		}

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Join the channel if it's valid
			if channelRegexp.MatchString(channel) {
				client.joinChannel(channel)
			}
		}

	case "PART":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		reason := strings.Join(args[1:], " ")

		channels := strings.Split(args[0], ",")
		for _, channel := range channels {
			//Part the channel if it's valid
			if channelRegexp.MatchString(channel) {
				client.partChannel(channel, reason)
			}
		}

	case "PRIVMSG":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		message := strings.Join(args[1:], " ")

		channel, chanExists := s.channelMap[strings.ToLower(args[0])]
		client2, clientExists := s.clientMap[strings.ToLower(args[0])]

		if chanExists {
			if channel.mode.noExternal {
				if _, inChannel := channel.clientMap[strings.ToLower(client.nick)]; !inChannel {
					//Not in channel, not allowed to send
					client.reply(errCannotSend, args[0])
					return
				}
			}
			if channel.mode.moderated {
				clientMode := channel.modeMap[strings.ToLower(client.nick)]
				if !clientMode.operator && !clientMode.voice {
					//It's moderated and we're not +v or +o, do nothing
					client.reply(errCannotSend, args[0])
					return
				}
			}
			for _, c := range channel.clientMap {
				if c != client {
					c.reply(rplMsg, client.nick, args[0], message)
				}
			}
		} else if clientExists {
			client.reply(rplMsg, client.nick, client2.nick, message)
		} else {
			client.reply(errNoSuchNick, args[0])
		}

	case "QUIT":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		client.disconnect()

	case "TOPIC":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		channel, exists := s.channelMap[strings.ToLower(args[0])]
		if exists == false {
			client.reply(errNoSuchNick, args[0])
			return
		}

		if len(args) == 1 {
			client.reply(rplTopic, channel.name, channel.topic)
			return
		}

		clientMode := channel.modeMap[strings.ToLower(client.nick)]
		if channel.mode.topicLocked && !clientMode.operator {
			client.reply(errNoPriv)
			return
		}

		if args[1] == ":" {
			channel.topic = ""
			for _, client := range channel.clientMap {
				client.reply(rplNoTopic, channel.name)
			}
		} else {
			topic := strings.Join(args[1:], " ")
			topic = strings.TrimPrefix(topic, ":")
			channel.topic = topic

			for _, client := range channel.clientMap {
				client.reply(rplTopic, channel.name, channel.topic)
			}
		}

	case "LIST":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) == 0 {
			chanList := make([]string, 0, len(s.channelMap))

			for channelName, channel := range s.channelMap {
				if channel.mode.secret {
					if _, inChannel := channel.clientMap[strings.ToLower(client.nick)]; !inChannel {
						//Not in the channel, skip
						continue
					}
				}
				listItem := fmt.Sprintf("%s %d :%s", channelName, len(channel.clientMap), channel.topic)
				chanList = append(chanList, listItem)
			}

			client.reply(rplList, chanList...)

		} else {
			channels := strings.Split(args[0], ",")
			chanList := make([]string, 0, len(channels))

			for _, channelName := range channels {
				if channel, exists := s.channelMap[strings.ToLower(channelName)]; exists {
					listItem := fmt.Sprintf("%s %d :%s", channelName, len(channel.clientMap), channel.topic)
					chanList = append(chanList, listItem)
				}
			}

			client.reply(rplList, chanList...)
		}
	case "OPER":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		username := args[0]
		password := args[1]

		if hashedPassword, exists := s.operatorMap[username]; exists {
			h := sha1.New()
			io.WriteString(h, password)
			pass := fmt.Sprintf("%x", h.Sum(nil))
			if hashedPassword == pass {
				client.operator = true
				client.reply(rplOper)
				return
			}
		}
		client.reply(errPassword)

	case "KILL":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if client.operator == false {
			client.reply(errNoPriv)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		nick := args[0]

		reason := strings.Join(args[1:], " ")

		client, exists := s.clientMap[strings.ToLower(nick)]
		if !exists {
			client.reply(errNoSuchNick, nick)
			return
		}

		client.reply(rplKill, client.nick, reason)
		client.disconnect()

	case "KICK":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 2 {
			client.reply(errMoreArgs)
			return
		}

		channelKey := strings.ToLower(args[0])
		targetKey := strings.ToLower(args[1])

		channel, channelExists := s.channelMap[channelKey]
		if !channelExists {
			client.reply(errNoSuchNick, args[0])
			return
		}

		target, targetExists := channel.clientMap[targetKey]
		if !targetExists {
			client.reply(errNoSuchNick, args[1])
			return
		}

		clientMode := channel.modeMap[client.key]
		if !clientMode.operator && !client.operator {
			client.reply(errNoPriv)
			return
		}

		reason := strings.Join(args[2:], " ")

		//It worked
		for _, client := range channel.clientMap {
			client.reply(rplKick, client.nick, channel.name, target.nick, reason)
		}

		delete(channel.clientMap, targetKey)
		delete(channel.modeMap, targetKey)
		delete(target.channelMap, channelKey)

	case "MODE":
		if client.registered == false {
			client.reply(errNotReg)
			return
		}

		if len(args) < 1 {
			client.reply(errMoreArgs)
			return
		}

		channelKey := strings.ToLower(args[0])

		channel, channelExists := s.channelMap[channelKey]
		if !channelExists {
			client.reply(errNoSuchNick, args[0])
			return
		}
		mode := channel.mode

		if len(args) == 1 {
			//No more args, they just want the mode
			client.reply(rplChannelModeIs, args[0], mode.String(), "")
			return
		}

		if cm, ok := channel.modeMap[strings.ToLower(client.nick)]; !ok || !cm.operator {
			//Not a channel operator.

			//If they're not an irc operator either, they'll fail
			if !client.operator {
				client.reply(errNoPriv)
				return
			}
		}

		hasClient := false
		var oldClientMode, newClientMode *ClientMode
		var targetClient *Client
		if len(args) >= 3 {
			clientKey := strings.ToLower(args[2])
			oldClientMode, hasClient = channel.modeMap[clientKey]
			if hasClient {
				targetClient = channel.clientMap[clientKey]
				newClientMode = new(ClientMode)
				*newClientMode = *oldClientMode
			}
		}

		mod := strings.ToLower(args[1])
		if strings.HasPrefix(mod, "+") {
			for _, char := range mod {
				switch char {
				case 's':
					mode.secret = true
				case 't':
					mode.topicLocked = true
				case 'm':
					mode.moderated = true
				case 'n':
					mode.noExternal = true
				case 'o':
					if hasClient {
						newClientMode.operator = true
					}
				case 'v':
					if hasClient {
						newClientMode.voice = true
					}
				}
			}
		} else if strings.HasPrefix(mod, "-") {
			for _, char := range mod {
				switch char {
				case 's':
					mode.secret = false
				case 't':
					mode.topicLocked = false
				case 'm':
					mode.moderated = false
				case 'n':
					mode.noExternal = false
				case 'o':
					if hasClient {
						newClientMode.operator = false
					}
				case 'v':
					if hasClient {
						newClientMode.voice = false
					}
				}
			}
		}

		if hasClient {
			*oldClientMode = *newClientMode
		}
		channel.mode = mode

		for _, client := range channel.clientMap {
			if hasClient {
				client.reply(rplChannelModeIs, channel.name, args[1], targetClient.nick)
			} else {
				client.reply(rplChannelModeIs, channel.name, args[1], "")
			}
		}

	default:
		client.reply(errUnknownCommand, command)
	}
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "log"
  "net/http"
  "strings"
  "time"
  "encoding/json"
  "flag"

  "github.com/marcelocajueiro/url_shortener/urls"
)

var (
  port    int
  urlBase string
)

func init() {
  flag.IntVar(&port, "p", 8888 , "port")
  urlBase = fmt.Sprintf("http://localhost:%d", port)

  flag.Parse()
}

func main() {
  stats := make (chan string)
  defer close (stats)
  go newStatistic(stats)

  http.HandleFunc("/api/shorten", Shortener)
  http.Handle("/r/", &Redirector{stats})
  http.HandleFunc("/api/stats/", StatsViewer)

  printLog("Starting server on port %d...", port)
  log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))
}

type Redirector struct {
  stats chan string
}
type Headers map[string]string
type Url struct {
  Id        string
  CreatedAt time.Time
  Destiny   string
}

func Shortener(w http.ResponseWriter, r *http.Request) {
  if r.Method != "POST" {
    respondWith(w, http.StatusMethodNotAllowed, Headers{
      "Allow": "POST",
    })
    return
  }

  url, new, err := urls.FindOrCreateNewUrl(extractUrl(r))

  if err != nil {
    respondWith(w, http.StatusBadRequest, nil)
    return
  }

  var status int

  if new {
    status = http.StatusCreated
  } else {
    status = http.StatusOK
  }

  shortUrl := fmt.Sprintf("%s/r/%s", urlBase, url.Id)
  respondWith(w, status, Headers{
    "Location": shortUrl,
    "Link": fmt.Sprintf("<%s/api/stats/%s>; rel=\"stats\"", urlBase, url.Id),
  })

  printLog("URL %s successfully shortened to %s", url.Destiny, shortUrl)
}

func (red *Redirector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  findUrlAndExecute(w, r, func (url *urls.Url) {
    http.Redirect(w, r, url.Destiny, http.StatusMovedPermanently)
    red.stats <- url.Id
  })
}

func StatsViewer(w http.ResponseWriter, r *http.Request) {
  findUrlAndExecute(w, r, func (url *urls.Url) {
    json, err := json.Marshal(url.Stats())

    if err != nil {
      w.WriteHeader(http.StatusInternalServerError)
      return
    }

    respondWithJSON(w, string(json))
  })
}

func findUrlAndExecute(
  w http.ResponseWriter,
  r *http.Request,
  executor func(*urls.Url),
) {
  path := strings.Split(r.URL.Path, "/")
  id := path[len(path) -1]

  if url := urls.Search(id); url != nil {
    executor(url)
  } else {
    http.NotFound(w, r)
  }
}

func respondWith(w http.ResponseWriter, status int, headers Headers) {
  for k, v := range headers {
    w.Header().Set(k, v)
  }
  w.WriteHeader(status)
}

func respondWithJSON(w http.ResponseWriter, response string) {
  respondWith(w, http.StatusOK, Headers{
    "Content-Type": "application/json",
  })
  fmt.Fprintf(w, response)
}

func extractUrl(r *http.Request) string {
  url := make([]byte, r.ContentLength, r.ContentLength)
  r.Body.Read(url)
  return string (url)
}

func newStatistic(ids <-chan string) {
  for id := range ids {
    urls.RegisterClick(id)
    printLog("%s was clicked", id)
  }
}

func printLog(format string, values ...interface{}) {
  // log.Printf("Testando %d %d %d", 1, 2, 3)
  log.Printf(fmt.Sprintf("%s\n", format), values...)
}
<commit_msg>Change the urls package path<commit_after>package main

import (
  "fmt"
  "log"
  "net/http"
  "strings"
  "time"
  "encoding/json"
  "flag"

  "github.com/marcelocajueiro/simple_url_shortener/urls"
)

var (
  port    int
  urlBase string
)

func init() {
  flag.IntVar(&port, "p", 8888 , "to set a custom port")
  urlBase = fmt.Sprintf("http://localhost:%d", port)

  flag.Parse()
}

func main() {
  stats := make (chan string)
  defer close (stats)
  go newStatistic(stats)

  http.HandleFunc("/api/shorten", Shortener)
  http.Handle("/r/", &Redirector{stats})
  http.HandleFunc("/api/stats/", StatsViewer)

  printLog("Starting server on port %d...", port)
  log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))
}

type Redirector struct {
  stats chan string
}
type Headers map[string]string
type Url struct {
  Id        string
  CreatedAt time.Time
  Destiny   string
}

func Shortener(w http.ResponseWriter, r *http.Request) {
  if r.Method != "POST" {
    respondWith(w, http.StatusMethodNotAllowed, Headers{
      "Allow": "POST",
    })
    return
  }

  url, new, err := urls.FindOrCreateNewUrl(extractUrl(r))

  if err != nil {
    respondWith(w, http.StatusBadRequest, nil)
    return
  }

  var status int

  if new {
    status = http.StatusCreated
  } else {
    status = http.StatusOK
  }

  shortUrl := fmt.Sprintf("%s/r/%s", urlBase, url.Id)
  respondWith(w, status, Headers{
    "Location": shortUrl,
    "Link": fmt.Sprintf("<%s/api/stats/%s>; rel=\"stats\"", urlBase, url.Id),
  })

  printLog("URL %s successfully shortened to %s", url.Destiny, shortUrl)
}

func (red *Redirector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  findUrlAndExecute(w, r, func (url *urls.Url) {
    http.Redirect(w, r, url.Destiny, http.StatusMovedPermanently)
    red.stats <- url.Id
  })
}

func StatsViewer(w http.ResponseWriter, r *http.Request) {
  findUrlAndExecute(w, r, func (url *urls.Url) {
    json, err := json.Marshal(url.Stats())

    if err != nil {
      w.WriteHeader(http.StatusInternalServerError)
      return
    }

    respondWithJSON(w, string(json))
  })
}

func findUrlAndExecute(
  w http.ResponseWriter,
  r *http.Request,
  executor func(*urls.Url),
) {
  path := strings.Split(r.URL.Path, "/")
  id := path[len(path) -1]

  if url := urls.Search(id); url != nil {
    executor(url)
  } else {
    http.NotFound(w, r)
  }
}

func respondWith(w http.ResponseWriter, status int, headers Headers) {
  for k, v := range headers {
    w.Header().Set(k, v)
  }
  w.WriteHeader(status)
}

func respondWithJSON(w http.ResponseWriter, response string) {
  respondWith(w, http.StatusOK, Headers{
    "Content-Type": "application/json",
  })
  fmt.Fprintf(w, response)
}

func extractUrl(r *http.Request) string {
  url := make([]byte, r.ContentLength, r.ContentLength)
  r.Body.Read(url)
  return string (url)
}

func newStatistic(ids <-chan string) {
  for id := range ids {
    urls.RegisterClick(id)
    printLog("%s was clicked", id)
  }
}

func printLog(format string, values ...interface{}) {
  // log.Printf("Testando %d %d %d", 1, 2, 3)
  log.Printf(fmt.Sprintf("%s\n", format), values...)
}
<|endoftext|>


<commit_before>// Package tunnel is a server/client package that enables to proxy public
// connections to your local machine over a tunnel connection from the local
// machine to the public server.
package tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/yamux"
	"github.com/koding/logging"
)

var (
	errNoClientSession = errors.New("no client session established")
)

// Server is responsible for proxying public connections to the client over a
// tunnel connection. It also listens to control messages from the client.
type Server struct {
	// pending contains the channel that is associated with each new tunnel request
	pending   map[string]chan net.Conn
	pendingMu sync.Mutex // protects the pending map

	// sessions contains a session per virtual host. Sessions provides
	// multiplexing over one connection
	sessions   map[string]*yamux.Session
	sessionsMu sync.Mutex // protects the sessions map

	// controls contains the control connection from the client to the server
	controls *controls

	// virtualHosts is used to map public hosts to remote clients
	virtualHosts vhostStorage

	// onDisconnect contains the onDisconnect for each map
	onDisconnect   map[string]func() error
	onDisconnectMu sync.Mutex // protects onDisconnects

	// yamuxConfig is passed to new yamux.Session's
	yamuxConfig *yamux.Config

	log logging.Logger
}

// ServerConfig defines the configuration for the Server
type ServerConfig struct {
	// Debug enables debug mode, enable only if you want to debug the server
	Debug bool

	// Log defines the logger. If nil a default logging.Logger is used.
	Log logging.Logger

	// YamuxConfig defines the config which passed to every new yamux.Session. If nil
	// yamux.DefaultConfig() is used.
	YamuxConfig *yamux.Config
}

// NewServer creates a new Server. The defaults are used if config is nil.
func NewServer(cfg *ServerConfig) (*Server, error) {
	yamuxConfig := yamux.DefaultConfig()
	if cfg.YamuxConfig != nil {
		if err := yamux.VerifyConfig(cfg.YamuxConfig); err != nil {
			return nil, err
		}

		yamuxConfig = cfg.YamuxConfig
	}

	log := newLogger("tunnel-server", cfg.Debug)
	if cfg.Log != nil {
		log = cfg.Log
	}

	return &Server{
		pending:      make(map[string]chan net.Conn),
		sessions:     make(map[string]*yamux.Session),
		onDisconnect: make(map[string]func() error),
		virtualHosts: newVirtualHosts(),
		controls:     newControls(),
		yamuxConfig:  yamuxConfig,
		log:          log,
	}, nil
}

// ServeHTTP is a tunnel that creates an http/websocket tunnel between a
// public connection and the client connection.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// if the user didn't add the control and tunnel handler manually, we'll
	// going to infer and call the respective path handlers.
	switch path.Clean(r.URL.Path) + "/" {
	case controlPath:
		s.checkConnect(s.controlHandler).ServeHTTP(w, r)
		return
	}

	if err := s.handleHTTP(w, r); err != nil {
		http.Error(w, err.Error(), 502)
		return
	}
}

// handleHTTP handles a single HTTP request
func (s *Server) handleHTTP(w http.ResponseWriter, r *http.Request) error {
	s.log.Debug("HandleHTTP request:")
	s.log.Debug("%v", r)

	host := strings.ToLower(r.Host)
	if host == "" {
		return errors.New("request host is empty")
	}

	// get the identifier associated with this host
	identifier, ok := s.getIdentifier(host)
	if !ok {
		return fmt.Errorf("no virtual host available for %s", host)
	}

	// then grab the control connection that is associated with this identifier
	control, ok := s.getControl(identifier)
	if !ok {
		return errNoClientSession
	}

	session, err := s.getSession(identifier)
	if err != nil {
		return err
	}

	// if someoone hits foo.example.com:8080, this should be proxied to
	// localhost:8080, so send the port to the client so it knows how to proxy
	// correctly. If no port is available, it's up to client how to intepret it
	_, port, _ := net.SplitHostPort(r.Host)
	msg := controlMsg{
		Action:    requestClientSession,
		Protocol:  httpTransport,
		LocalPort: port,
	}

	s.log.Debug("Sending control msg %+v", msg)

	// ask client to open a session to us, so we can accept it
	if err := control.send(msg); err != nil {
		if err == errControlClosed || err == yamux.ErrStreamClosed {
			return errNoClientSession
		}

		return err
	}

	var stream net.Conn
	defer func() {
		if stream != nil {
			s.log.Debug("Closing stream")
			stream.Close()
		}
	}()

	acceptStream := func() error {
		stream, err = session.Accept()
		return err
	}

	// if we don't receive anything from the client, we'll timeout
	s.log.Debug("Waiting for session accept")
	select {
	case err := <-async(acceptStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		return errors.New("timeout getting session")
	}

	if err := r.Write(stream); err != nil {
		return err
	}

	s.log.Debug("Session opened to client, writing request to client")
	resp, err := http.ReadResponse(bufio.NewReader(stream), r)
	if err != nil {
		return fmt.Errorf("read from tunnel: %s", err.Error())
	}

	defer func() {
		if resp.Body != nil {
			if err := resp.Body.Close(); err != nil {
				s.log.Error("resp.Body Close error: %s", err.Error())
			}
		}
	}()

	s.log.Debug("Response received, writing back to public connection")
	s.log.Debug("%+v", resp)

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)

	if _, err := io.Copy(w, resp.Body); err != nil {
		s.log.Error("copy err: %s", err) // do not return, because we might write multipe headers
	}

	s.log.Debug("Response copy is finished")
	return nil
}

// controlHandler is used to capture incoming tunnel connect requests into raw
// tunnel TCP connections.
func (s *Server) controlHandler(w http.ResponseWriter, r *http.Request) (ctErr error) {
	identifier := r.Header.Get(xKTunnelIdentifier)
	_, ok := s.getHost(identifier)
	if !ok {
		return fmt.Errorf("no host associated for identifier %s. please use server.AddHost()", identifier)
	}

	ct, ok := s.getControl(identifier)
	if ok {
		ct.Close()
		s.log.Warning("Control connection for '%s' already exists. This is a race condition and needs to be fixed on client implementation", identifier)
		return fmt.Errorf("control conn for %s already exist. \n", identifier)
	}

	s.log.Debug("Tunnel with identifier %s", identifier)

	hj, ok := w.(http.Hijacker)
	if !ok {
		return errors.New("webserver doesn't support hijacking")
	}

	conn, _, err := hj.Hijack()
	if err != nil {
		return fmt.Errorf("hijack not possible %s", err)
	}

	io.WriteString(conn, "HTTP/1.1 "+connected+"\n\n")

	conn.SetDeadline(time.Time{})

	s.log.Debug("Creating control session")
	session, err := yamux.Server(conn, s.yamuxConfig)
	if err != nil {
		return err
	}
	s.addSession(identifier, session)

	var stream net.Conn

	// close and delete the session/stream if something goes wrong
	defer func() {
		if ctErr != nil {
			if stream != nil {
				stream.Close()
			}
			s.deleteSession(identifier)
		}
	}()

	acceptStream := func() error {
		stream, err = session.Accept()
		return err
	}

	// if we don't receive anything from the client, we'll timeout
	select {
	case err := <-async(acceptStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		return errors.New("timeout getting session")
	}

	s.log.Debug("Initiating handshake protocol")
	buf := make([]byte, len(ctHandshakeRequest))
	if _, err := stream.Read(buf); err != nil {
		return err
	}

	if string(buf) != ctHandshakeRequest {
		return fmt.Errorf("handshake aborted. got: %s", string(buf))
	}

	if _, err := stream.Write([]byte(ctHandshakeResponse)); err != nil {
		return err
	}

	// setup control stream and start to listen to messages
	ct = newControl(stream)
	s.addControl(identifier, ct)
	go s.listenControl(ct)

	s.log.Debug("Control connection is setup")
	return nil
}

// listenControl listens to messages coming from the client.
func (s *Server) listenControl(ct *control) {
	for {
		var msg map[string]interface{}
		err := ct.dec.Decode(&msg)
		if err != nil {
			ct.Close()
			s.deleteControl(ct.identifier)
			s.deleteSession(ct.identifier)
			if err := s.callOnDisconect(ct.identifier); err != nil {
				s.log.Error("onDisconnect (%s) err: %s", ct.identifier, err)
			}

			if err != io.EOF {
				s.log.Error("decode err: %s", err)
			}
			return
		}

		// right now we don't do anything with the messages, but because the
		// underlying connection needs to establihsed, we know when we have
		// disconnection(above), so we can cleanup the connection.
		s.log.Debug("msg: %s", msg)
	}
}

// OnDisconnect calls the function when the client connected with the
// associated identifier disconnects from the server. After a client is
// disconnected, the associated function is alro removed and needs to be
// readded again.
func (s *Server) OnDisconnect(identifier string, fn func() error) {
	s.onDisconnectMu.Lock()
	s.onDisconnect[identifier] = fn
	s.onDisconnectMu.Unlock()
}

func (s *Server) callOnDisconect(identifier string) error {
	s.onDisconnectMu.Lock()
	defer s.onDisconnectMu.Unlock()

	fn, ok := s.onDisconnect[identifier]
	if !ok {
		return nil
	}

	// delete after we are finished with it
	delete(s.onDisconnect, identifier)

	if fn == nil {
		return errors.New("onDisconnect function for '%s' is set to nil")
	}

	return fn()
}

// AddHost adds the given virtual host and maps it to the identifier.
func (s *Server) AddHost(host, identifier string) {
	s.virtualHosts.AddHost(host, identifier)
}

// DeleteHost deletes the given virtual host. Once removed any request to this
// host is denied.
func (s *Server) DeleteHost(host string) {
	s.virtualHosts.DeleteHost(host)
}

func (s *Server) getIdentifier(host string) (string, bool) {
	identifier, ok := s.virtualHosts.GetIdentifier(host)
	return identifier, ok
}

func (s *Server) getHost(identifier string) (string, bool) {
	host, ok := s.virtualHosts.GetHost(identifier)
	return host, ok
}

func (s *Server) addControl(identifier string, conn *control) {
	s.controls.addControl(identifier, conn)
}

func (s *Server) getControl(identifier string) (*control, bool) {
	return s.controls.getControl(identifier)
}

func (s *Server) deleteControl(identifier string) {
	s.controls.deleteControl(identifier)
}

func (s *Server) getSession(identifier string) (*yamux.Session, error) {
	s.sessionsMu.Lock()
	session, ok := s.sessions[identifier]
	s.sessionsMu.Unlock()

	if !ok {
		return nil, fmt.Errorf("no session available for identifier: '%s'", identifier)
	}

	return session, nil
}

func (s *Server) addSession(identifier string, session *yamux.Session) {
	s.sessionsMu.Lock()
	s.sessions[identifier] = session
	s.sessionsMu.Unlock()
}

func (s *Server) deleteSession(identifier string) {
	s.sessionsMu.Lock()
	session, ok := s.sessions[identifier]
	s.sessionsMu.Unlock()

	if !ok {
		return // nothing to delete
	}

	if session != nil {
		session.GoAway() // don't accept any new connection
		session.Close()
	}

	delete(s.sessions, identifier)
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// checkConnect checks wether the incoming request is HTTP CONNECT method. If
func (s *Server) checkConnect(fn func(w http.ResponseWriter, r *http.Request) error) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "CONNECT" {
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(http.StatusMethodNotAllowed)
			io.WriteString(w, "405 must CONNECT\n")
			return
		}

		err := fn(w, r)
		if err != nil {
			s.log.Error("Handler err: %v", err.Error())
			http.Error(w, err.Error(), 502)
		}
	})
}

func newLogger(name string, debug bool) logging.Logger {
	log := logging.NewLogger(name)
	logHandler := logging.NewWriterHandler(os.Stderr)
	logHandler.Colorize = true
	log.SetHandler(logHandler)

	if debug {
		log.SetLevel(logging.DEBUG)
		logHandler.SetLevel(logging.DEBUG)
	}

	return log
}
<commit_msg>server: handle yamux.ErrSessionShutdown too<commit_after>// Package tunnel is a server/client package that enables to proxy public
// connections to your local machine over a tunnel connection from the local
// machine to the public server.
package tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/yamux"
	"github.com/koding/logging"
)

var (
	errNoClientSession = errors.New("no client session established")
)

// Server is responsible for proxying public connections to the client over a
// tunnel connection. It also listens to control messages from the client.
type Server struct {
	// pending contains the channel that is associated with each new tunnel request
	pending   map[string]chan net.Conn
	pendingMu sync.Mutex // protects the pending map

	// sessions contains a session per virtual host. Sessions provides
	// multiplexing over one connection
	sessions   map[string]*yamux.Session
	sessionsMu sync.Mutex // protects the sessions map

	// controls contains the control connection from the client to the server
	controls *controls

	// virtualHosts is used to map public hosts to remote clients
	virtualHosts vhostStorage

	// onDisconnect contains the onDisconnect for each map
	onDisconnect   map[string]func() error
	onDisconnectMu sync.Mutex // protects onDisconnects

	// yamuxConfig is passed to new yamux.Session's
	yamuxConfig *yamux.Config

	log logging.Logger
}

// ServerConfig defines the configuration for the Server
type ServerConfig struct {
	// Debug enables debug mode, enable only if you want to debug the server
	Debug bool

	// Log defines the logger. If nil a default logging.Logger is used.
	Log logging.Logger

	// YamuxConfig defines the config which passed to every new yamux.Session. If nil
	// yamux.DefaultConfig() is used.
	YamuxConfig *yamux.Config
}

// NewServer creates a new Server. The defaults are used if config is nil.
func NewServer(cfg *ServerConfig) (*Server, error) {
	yamuxConfig := yamux.DefaultConfig()
	if cfg.YamuxConfig != nil {
		if err := yamux.VerifyConfig(cfg.YamuxConfig); err != nil {
			return nil, err
		}

		yamuxConfig = cfg.YamuxConfig
	}

	log := newLogger("tunnel-server", cfg.Debug)
	if cfg.Log != nil {
		log = cfg.Log
	}

	return &Server{
		pending:      make(map[string]chan net.Conn),
		sessions:     make(map[string]*yamux.Session),
		onDisconnect: make(map[string]func() error),
		virtualHosts: newVirtualHosts(),
		controls:     newControls(),
		yamuxConfig:  yamuxConfig,
		log:          log,
	}, nil
}

// ServeHTTP is a tunnel that creates an http/websocket tunnel between a
// public connection and the client connection.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// if the user didn't add the control and tunnel handler manually, we'll
	// going to infer and call the respective path handlers.
	switch path.Clean(r.URL.Path) + "/" {
	case controlPath:
		s.checkConnect(s.controlHandler).ServeHTTP(w, r)
		return
	}

	if err := s.handleHTTP(w, r); err != nil {
		http.Error(w, err.Error(), 502)
		return
	}
}

// handleHTTP handles a single HTTP request
func (s *Server) handleHTTP(w http.ResponseWriter, r *http.Request) error {
	s.log.Debug("HandleHTTP request:")
	s.log.Debug("%v", r)

	host := strings.ToLower(r.Host)
	if host == "" {
		return errors.New("request host is empty")
	}

	// get the identifier associated with this host
	identifier, ok := s.getIdentifier(host)
	if !ok {
		return fmt.Errorf("no virtual host available for %s", host)
	}

	// then grab the control connection that is associated with this identifier
	control, ok := s.getControl(identifier)
	if !ok {
		return errNoClientSession
	}

	session, err := s.getSession(identifier)
	if err != nil {
		return err
	}

	// if someoone hits foo.example.com:8080, this should be proxied to
	// localhost:8080, so send the port to the client so it knows how to proxy
	// correctly. If no port is available, it's up to client how to intepret it
	_, port, _ := net.SplitHostPort(r.Host)
	msg := controlMsg{
		Action:    requestClientSession,
		Protocol:  httpTransport,
		LocalPort: port,
	}

	s.log.Debug("Sending control msg %+v", msg)

	// ask client to open a session to us, so we can accept it
	if err := control.send(msg); err != nil {
		switch err {
		case errControlClosed, yamux.ErrStreamClosed, yamux.ErrSessionShutdown:
			return errNoClientSession
		default:
			return err
		}
	}

	var stream net.Conn
	defer func() {
		if stream != nil {
			s.log.Debug("Closing stream")
			stream.Close()
		}
	}()

	acceptStream := func() error {
		stream, err = session.Accept()
		return err
	}

	// if we don't receive anything from the client, we'll timeout
	s.log.Debug("Waiting for session accept")
	select {
	case err := <-async(acceptStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		return errors.New("timeout getting session")
	}

	if err := r.Write(stream); err != nil {
		return err
	}

	s.log.Debug("Session opened to client, writing request to client")
	resp, err := http.ReadResponse(bufio.NewReader(stream), r)
	if err != nil {
		return fmt.Errorf("read from tunnel: %s", err.Error())
	}

	defer func() {
		if resp.Body != nil {
			if err := resp.Body.Close(); err != nil {
				s.log.Error("resp.Body Close error: %s", err.Error())
			}
		}
	}()

	s.log.Debug("Response received, writing back to public connection")
	s.log.Debug("%+v", resp)

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)

	if _, err := io.Copy(w, resp.Body); err != nil {
		s.log.Error("copy err: %s", err) // do not return, because we might write multipe headers
	}

	s.log.Debug("Response copy is finished")
	return nil
}

// controlHandler is used to capture incoming tunnel connect requests into raw
// tunnel TCP connections.
func (s *Server) controlHandler(w http.ResponseWriter, r *http.Request) (ctErr error) {
	identifier := r.Header.Get(xKTunnelIdentifier)
	_, ok := s.getHost(identifier)
	if !ok {
		return fmt.Errorf("no host associated for identifier %s. please use server.AddHost()", identifier)
	}

	ct, ok := s.getControl(identifier)
	if ok {
		ct.Close()
		s.log.Warning("Control connection for '%s' already exists. This is a race condition and needs to be fixed on client implementation", identifier)
		return fmt.Errorf("control conn for %s already exist. \n", identifier)
	}

	s.log.Debug("Tunnel with identifier %s", identifier)

	hj, ok := w.(http.Hijacker)
	if !ok {
		return errors.New("webserver doesn't support hijacking")
	}

	conn, _, err := hj.Hijack()
	if err != nil {
		return fmt.Errorf("hijack not possible %s", err)
	}

	io.WriteString(conn, "HTTP/1.1 "+connected+"\n\n")

	conn.SetDeadline(time.Time{})

	s.log.Debug("Creating control session")
	session, err := yamux.Server(conn, s.yamuxConfig)
	if err != nil {
		return err
	}
	s.addSession(identifier, session)

	var stream net.Conn

	// close and delete the session/stream if something goes wrong
	defer func() {
		if ctErr != nil {
			if stream != nil {
				stream.Close()
			}
			s.deleteSession(identifier)
		}
	}()

	acceptStream := func() error {
		stream, err = session.Accept()
		return err
	}

	// if we don't receive anything from the client, we'll timeout
	select {
	case err := <-async(acceptStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		return errors.New("timeout getting session")
	}

	s.log.Debug("Initiating handshake protocol")
	buf := make([]byte, len(ctHandshakeRequest))
	if _, err := stream.Read(buf); err != nil {
		return err
	}

	if string(buf) != ctHandshakeRequest {
		return fmt.Errorf("handshake aborted. got: %s", string(buf))
	}

	if _, err := stream.Write([]byte(ctHandshakeResponse)); err != nil {
		return err
	}

	// setup control stream and start to listen to messages
	ct = newControl(stream)
	s.addControl(identifier, ct)
	go s.listenControl(ct)

	s.log.Debug("Control connection is setup")
	return nil
}

// listenControl listens to messages coming from the client.
func (s *Server) listenControl(ct *control) {
	for {
		var msg map[string]interface{}
		err := ct.dec.Decode(&msg)
		if err != nil {
			ct.Close()
			s.deleteControl(ct.identifier)
			s.deleteSession(ct.identifier)
			if err := s.callOnDisconect(ct.identifier); err != nil {
				s.log.Error("onDisconnect (%s) err: %s", ct.identifier, err)
			}

			if err != io.EOF {
				s.log.Error("decode err: %s", err)
			}
			return
		}

		// right now we don't do anything with the messages, but because the
		// underlying connection needs to establihsed, we know when we have
		// disconnection(above), so we can cleanup the connection.
		s.log.Debug("msg: %s", msg)
	}
}

// OnDisconnect calls the function when the client connected with the
// associated identifier disconnects from the server. After a client is
// disconnected, the associated function is alro removed and needs to be
// readded again.
func (s *Server) OnDisconnect(identifier string, fn func() error) {
	s.onDisconnectMu.Lock()
	s.onDisconnect[identifier] = fn
	s.onDisconnectMu.Unlock()
}

func (s *Server) callOnDisconect(identifier string) error {
	s.onDisconnectMu.Lock()
	defer s.onDisconnectMu.Unlock()

	fn, ok := s.onDisconnect[identifier]
	if !ok {
		return nil
	}

	// delete after we are finished with it
	delete(s.onDisconnect, identifier)

	if fn == nil {
		return errors.New("onDisconnect function for '%s' is set to nil")
	}

	return fn()
}

// AddHost adds the given virtual host and maps it to the identifier.
func (s *Server) AddHost(host, identifier string) {
	s.virtualHosts.AddHost(host, identifier)
}

// DeleteHost deletes the given virtual host. Once removed any request to this
// host is denied.
func (s *Server) DeleteHost(host string) {
	s.virtualHosts.DeleteHost(host)
}

func (s *Server) getIdentifier(host string) (string, bool) {
	identifier, ok := s.virtualHosts.GetIdentifier(host)
	return identifier, ok
}

func (s *Server) getHost(identifier string) (string, bool) {
	host, ok := s.virtualHosts.GetHost(identifier)
	return host, ok
}

func (s *Server) addControl(identifier string, conn *control) {
	s.controls.addControl(identifier, conn)
}

func (s *Server) getControl(identifier string) (*control, bool) {
	return s.controls.getControl(identifier)
}

func (s *Server) deleteControl(identifier string) {
	s.controls.deleteControl(identifier)
}

func (s *Server) getSession(identifier string) (*yamux.Session, error) {
	s.sessionsMu.Lock()
	session, ok := s.sessions[identifier]
	s.sessionsMu.Unlock()

	if !ok {
		return nil, fmt.Errorf("no session available for identifier: '%s'", identifier)
	}

	return session, nil
}

func (s *Server) addSession(identifier string, session *yamux.Session) {
	s.sessionsMu.Lock()
	s.sessions[identifier] = session
	s.sessionsMu.Unlock()
}

func (s *Server) deleteSession(identifier string) {
	s.sessionsMu.Lock()
	session, ok := s.sessions[identifier]
	s.sessionsMu.Unlock()

	if !ok {
		return // nothing to delete
	}

	if session != nil {
		session.GoAway() // don't accept any new connection
		session.Close()
	}

	delete(s.sessions, identifier)
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// checkConnect checks wether the incoming request is HTTP CONNECT method. If
func (s *Server) checkConnect(fn func(w http.ResponseWriter, r *http.Request) error) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "CONNECT" {
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(http.StatusMethodNotAllowed)
			io.WriteString(w, "405 must CONNECT\n")
			return
		}

		err := fn(w, r)
		if err != nil {
			s.log.Error("Handler err: %v", err.Error())
			http.Error(w, err.Error(), 502)
		}
	})
}

func newLogger(name string, debug bool) logging.Logger {
	log := logging.NewLogger(name)
	logHandler := logging.NewWriterHandler(os.Stderr)
	logHandler.Colorize = true
	log.SetHandler(logHandler)

	if debug {
		log.SetLevel(logging.DEBUG)
		logHandler.SetLevel(logging.DEBUG)
	}

	return log
}
<|endoftext|>


<commit_before>// Copyright 2016 Jacob Taylor jacob.taylor@gmail.com
// License: Apache2
package main

import (
    "fmt"
    "net"
    "./utils"
    "bufio"
    "encoding/binary"
    "time"
    "os"
)

func send_export_list_item(output *bufio.Writer, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer) {
    send_message(output, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    // attempt to open the file read only
    file, err := os.Open("/Users/jacob/work/nbd/sample_disks/happyu")
    utils.ErrorCheck(err)
    //defer file.Close()

    data := make([]byte, 1024)
    count, err := file.Read(data)
    utils.ErrorCheck(err)

    if count > 100 {
        count = 100
    }

    // send export information
    // size u64
    // flags u16
    // Zeros (124 bytes)?
    buffer := make([]byte, 256)
    offset := 0

    binary.BigEndian.PutUint64(buffer[offset:], 52428800)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 0)  // flags
    offset += 2

    //offset += 124       // zero pad

    len, err := output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)
    fmt.Printf("Wrote %d chars: %v\n", len, buffer[:offset])

    fmt.Printf("File descriptor:\n%+v\n", *file)
    fmt.Printf("First 100 bytes: \n%v\n", data[:count])

    // send a reply with the handle
    //S: 32 bits, 0x67446698, magic (NBD_REPLY_MAGIC)
    //S: 32 bits, error (MAY be zero)
    //S: 64 bits, handle
    //S: (length bytes of data if the request is of type NBD_CMD_READ)
    //offset = 0
    //binary.BigEndian.PutUint32(buffer[offset:], utils.NBD_REPLY_MAGIC)
    //offset += 4
    //
    //binary.BigEndian.PutUint32(buffer[offset:], 0) // error
    //offset += 4
    //
    //binary.BigEndian.PutUint64(buffer[offset:], 8000) // handle
    //offset += 8
    //
    //fmt.Printf("Writing out data: %v\n", buffer[:offset])
    //len, err = output.Write(buffer[:offset])
    //output.Flush()
    //utils.ErrorCheck(err)
    fmt.Printf("Done sending data\n")

    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Fetch the data until we get the initial options
    time.Sleep(300 * time.Millisecond)

    fmt.Printf("about to read\n")
    for ; ;  {
        var zero_time time.Time
        conn.SetReadDeadline(zero_time)
        short_data := make([]byte, 1)
        conn.Read(short_data)
        fmt.Printf("read byte: %v\n", short_data)
        time.Sleep(300 * time.Millisecond)

    }

    //conn2, err = listener.Accept()
    //utils.ErrorCheck(err)

    data = make([]byte, 1024)
    offset = 0
    waiting_for := 16       // wait for at least the minimum payload size

    for offset < waiting_for {
        fmt.Printf("1: offset: %d, data: %v\n", offset, data)
        length, err := conn.Read(data[offset:])
        offset += length
        fmt.Printf("3: offset: %d, err: %v, data: %v\n", offset, err, data)
        //utils.ErrorCheck(err)
        fmt.Printf("4: offset: %d, data: %v\n", offset, data)
        utils.LogData("Reading instruction", offset, data)
        fmt.Printf("5: offset: %d, data: %v\n", offset, data)
        if offset < waiting_for {
        fmt.Printf("6: offset: %d, data: %v\n", offset, data)
            time.Sleep(1000 * time.Millisecond)
        }
    }
    fmt.Printf("done reading\n")
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function

}

func send_export_list(output *bufio.Writer) {
    export_name_list := []string{"happy_export", "very_happy_export", "third_export"}

    for index := range export_name_list {
        send_export_list_item(output, export_name_list[index])
    }

    send_ack(output)
}

func send_message(output *bufio.Writer, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], uint32(3))  // Flags (3 = supports list)
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

func main() {
    controlC := [...]byte{255, 244, 255, 253, 6}
    listener, err := net.Listen("tcp", "192.168.214.1:8000")
    utils.ErrorCheck(err)

    fmt.Printf("Hello World, we have %v\n", listener)
    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.Flush()
        output.WriteString("IHAVEOPT")      // Magic
        output.Flush()
        output.Write([]byte{0, 3})          // Flags (3 = supports list)
        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }

        // Skip the first 8 characters (options)
        command := binary.BigEndian.Uint32(data[12:])
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        fmt.Sprintf("command is: %d\npayload_size is: %d\n", command, payload_size)
        waiting_for += int(payload_size)
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }
        payload := make([]byte, payload_size)

        if payload_size > 0{
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)
        fmt.Printf("command is: %v\n", command)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output)
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            export_name(output, conn, payload_size, payload)
            break
        }

        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function
        // Fetch the data until we get the initial options
        fmt.Printf("about to read\n")

        //conn2, err = listener.Accept()
        //utils.ErrorCheck(err)

        data = make([]byte, 1024)
        offset = 0
        waiting_for = 16       // wait for at least the minimum payload size

        for offset < waiting_for {
            fmt.Printf("1: offset: %d, data: %v\n", offset, data)
            length, err := conn.Read(data[offset:])
            offset += length
            fmt.Printf("3: offset: %d, err: %v, data: %v\n", offset, err, data)
            //utils.ErrorCheck(err)
            fmt.Printf("4: offset: %d, data: %v\n", offset, data)
            utils.LogData("Reading instruction", offset, data)
            fmt.Printf("5: offset: %d, data: %v\n", offset, data)
            if offset < waiting_for {
            fmt.Printf("6: offset: %d, data: %v\n", offset, data)
                time.Sleep(1000 * time.Millisecond)
            }
        }
        fmt.Printf("done reading\n")
        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function
        // Duplicated code. move this to helper function

        input := bufio.NewScanner(conn)
        for input.Scan() {
            if len(input.Bytes()) == 5 {
                temp := [...]byte{0, 0, 0, 0, 0}
                copy(temp[:], input.Bytes())
                if temp == controlC {
                    fmt.Printf("Control-C received. Bye\n")
                    break
                }
            }
            fmt.Printf("%s echo: %s '%v'\n", conn.RemoteAddr().String(), input.Text(), input.Bytes())
        }
        conn.Close()

    }

}
<commit_msg>adding loop for handling read write and disconnect requests<commit_after>// Copyright 2016 Jacob Taylor jacob.taylor@gmail.com
// License: Apache2
package main

import (
    "fmt"
    "net"
    "./utils"
    "bufio"
    "encoding/binary"
    "time"
    "os"
)

func send_export_list_item(output *bufio.Writer, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer) {
    send_message(output, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    // attempt to open the file read only
    file, err := os.Open("/Users/jacob/work/nbd/sample_disks/happyu")
    utils.ErrorCheck(err)
    //defer file.Close()

    data := make([]byte, 1024)
    count, err := file.Read(data)
    utils.ErrorCheck(err)

    if count > 100 {
        count = 100
    }

    // send export information
    // size u64
    // flags u16
    // Zeros (124 bytes)?
    buffer := make([]byte, 256)
    offset := 0

    binary.BigEndian.PutUint64(buffer[offset:], 52428800)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 0)  // flags
    offset += 2

    //offset += 124       // zero pad

    len, err := output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)
    fmt.Printf("Wrote %d chars: %v\n", len, buffer[:offset])

    fmt.Printf("File descriptor:\n%+v\n", *file)
    fmt.Printf("First 100 bytes: \n%v\n", data[:count])

    // send a reply with the handle
    //S: 32 bits, 0x67446698, magic (NBD_REPLY_MAGIC)
    //S: 32 bits, error (MAY be zero)
    //S: 64 bits, handle
    //S: (length bytes of data if the request is of type NBD_CMD_READ)
    //offset = 0
    //binary.BigEndian.PutUint32(buffer[offset:], utils.NBD_REPLY_MAGIC)
    //offset += 4
    //
    //binary.BigEndian.PutUint32(buffer[offset:], 0) // error
    //offset += 4
    //
    //binary.BigEndian.PutUint64(buffer[offset:], 8000) // handle
    //offset += 8
    //
    //fmt.Printf("Writing out data: %v\n", buffer[:offset])
    //len, err = output.Write(buffer[:offset])
    //output.Flush()
    //utils.ErrorCheck(err)
    fmt.Printf("Done sending data\n")

    buffer = make([]byte, 512*1024)
    for {

        offset := 0
        waiting_for := 24       // wait for at least the minimum payload size

        for offset < waiting_for {
            length, err := conn.Read(buffer[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, buffer)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }

        //magic := binary.BigEndian.Uint32(buffer)
        command := binary.BigEndian.Uint32(buffer[4:8])
        handle := binary.BigEndian.Uint64(buffer[8:16])
        from := binary.BigEndian.Uint64(buffer[16:24])
        length := binary.BigEndian.Uint64(buffer[16:24])

        switch command {
        case utils.NBD_COMMAND_READ:
            fmt.Printf("We have a request to read handle: %v, from: %v, length: %v\n", handle, from, length)
            continue
        case utils.NBD_COMMAND_WRITE:
            fmt.Printf("We have a request to write handle: %v, from: %v, length: %v\n", handle, from, length)
            continue
        case utils.NBD_COMMAND_DISCONNECT:
            fmt.Printf("We have received a request to disconnect\n")
            // close the file and return
            return
        }
    }

    //
    //
    //// Fetch the data until we get the initial options
    //data := make([]byte, 1024)
    //offset := 0
    //waiting_for := 16       // wait for at least the minimum payload size
    //
    //for offset < waiting_for {
    //    length, err := conn.Read(data[offset:])
    //    offset += length
    //    utils.ErrorCheck(err)
    //    utils.LogData("Reading instruction", offset, data)
    //    if offset < waiting_for {
    //        time.Sleep(5 * time.Millisecond)
    //    }
    //}
    //
    //// Skip the first 8 characters (options)
    //command := binary.BigEndian.Uint32(data[12:])
    //payload_size := int(binary.BigEndian.Uint32(data[16:]))
    //
    //
    //    syscall.Read(nbd.socket, buf[0:28])
    //
		//x.magic = binary.BigEndian.Uint32(buf)
		//x.typus = binary.BigEndian.Uint32(buf[4:8])
		//x.handle = binary.BigEndian.Uint64(buf[8:16])
		//x.from = binary.BigEndian.Uint64(buf[16:24])
		//x.len = binary.BigEndian.Uint32(buf[24:28])
    //
    //
    //
    //// Duplicated code. move this to helper function
    //// Duplicated code. move this to helper function
    //// Duplicated code. move this to helper function
    //// Duplicated code. move this to helper function
    //// Fetch the data until we get the initial options
    //time.Sleep(300 * time.Millisecond)
    //
    //fmt.Printf("about to read\n")
    //for ; ;  {
    //    var zero_time time.Time
    //    conn.SetReadDeadline(zero_time)
    //    short_data := make([]byte, 1)
    //    conn.Read(short_data)
    //    fmt.Printf("read byte: %v\n", short_data)
    //    time.Sleep(300 * time.Millisecond)
    //
    //}
    //
    //
    ////conn2, err = listener.Accept()
    ////utils.ErrorCheck(err)
    //
    //data = make([]byte, 1024)
    //offset = 0
    //waiting_for := 16       // wait for at least the minimum payload size
    //
    //for offset < waiting_for {
    //    fmt.Printf("1: offset: %d, data: %v\n", offset, data)
    //    length, err := conn.Read(data[offset:])
    //    offset += length
    //    fmt.Printf("3: offset: %d, err: %v, data: %v\n", offset, err, data)
    //    //utils.ErrorCheck(err)
    //    fmt.Printf("4: offset: %d, data: %v\n", offset, data)
    //    utils.LogData("Reading instruction", offset, data)
    //    fmt.Printf("5: offset: %d, data: %v\n", offset, data)
    //    if offset < waiting_for {
    //    fmt.Printf("6: offset: %d, data: %v\n", offset, data)
    //        time.Sleep(1000 * time.Millisecond)
    //    }
    //}
    fmt.Printf("done reading\n")
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function
    // Duplicated code. move this to helper function

}

func send_export_list(output *bufio.Writer) {
    export_name_list := []string{"happy_export", "very_happy_export", "third_export"}

    for index := range export_name_list {
        send_export_list_item(output, export_name_list[index])
    }

    send_ack(output)
}

func send_message(output *bufio.Writer, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], uint32(3))  // Flags (3 = supports list)
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

func main() {
    listener, err := net.Listen("tcp", "192.168.214.1:8000")
    utils.ErrorCheck(err)

    fmt.Printf("Hello World, we have %v\n", listener)
    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.Flush()
        output.WriteString("IHAVEOPT")      // Magic
        output.Flush()
        output.Write([]byte{0, 3})          // Flags (3 = supports list)
        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }

        // Skip the first 8 characters (options)
        command := binary.BigEndian.Uint32(data[12:])
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        fmt.Sprintf("command is: %d\npayload_size is: %d\n", command, payload_size)
        waiting_for += int(payload_size)
        for offset < waiting_for {
            length, err := conn.Read(data[offset:])
            offset += length
            utils.ErrorCheck(err)
            utils.LogData("Reading instruction", offset, data)
            if offset < waiting_for {
                time.Sleep(5 * time.Millisecond)
            }
        }
        payload := make([]byte, payload_size)

        if payload_size > 0{
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)
        fmt.Printf("command is: %v\n", command)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload)
            break
        }
    }

}
<|endoftext|>


<commit_before>// Package gotest is a standard Go test output parser.
package gotest

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
)

var (
	// regexBenchInfo captures 3-5 groups: benchmark name, number of times ran, ns/op (with or without decimal), MB/sec (optional), B/op (optional), and allocs/op (optional).
	regexBenchmark    = regexp.MustCompile(`^(Benchmark[^ -]+)$`)
	regexBenchSummary = regexp.MustCompile(`^(Benchmark[^ -]+)(?:-\d+\s+|\s+)(\d+)\s+(\d+|\d+\.\d+)\sns\/op(?:\s+(\d+|\d+\.\d+)\sMB\/s)?(?:\s+(\d+)\sB\/op)?(?:\s+(\d+)\sallocs/op)?`)
	regexCoverage     = regexp.MustCompile(`^coverage:\s+(\d+|\d+\.\d+)%\s+of\s+statements(?:\sin\s(.+))?$`)
	regexEndBenchmark = regexp.MustCompile(`^--- (BENCH|FAIL|SKIP): (Benchmark[^ -]+)(?:-\d+)?$`)
	regexEndTest      = regexp.MustCompile(`((?:    )*)--- (PASS|FAIL|SKIP): ([^ ]+) \((\d+\.\d+)(?: seconds|s)\)`)
	regexStatus       = regexp.MustCompile(`^(PASS|FAIL|SKIP)$`)
	regexSummary      = regexp.MustCompile(`` +
		// 1: result
		`^(\?|ok|FAIL)` +
		// 2: package name
		`\s+([^ \t]+)` +
		// 3: duration (optional)
		`(?:\s+(\d+\.\d+)s)?` +
		// 4: cached indicator (optional)
		`(?:\s+(\(cached\)))?` +
		// 5: [status message] (optional)
		`(?:\s+(\[[^\]]+\]))?` +
		// 6: coverage percentage (optional)
		// 7: coverage package list (optional)
		`(?:\s+coverage:\s+(\d+\.\d+)%\sof\sstatements(?:\sin\s(.+))?)?$`)
)

// Option defines options that can be passed to gotest.New.
type Option func(*Parser)

// PackageName is an Option that sets the default package name to use when it
// cannot be determined from the test output.
func PackageName(name string) Option {
	return func(p *Parser) {
		p.packageName = name
	}
}

// TimestampFunc is an Option that sets the timestamp function that is used to
// determine the current time when creating the Report. This can be used to
// override the default behaviour of using time.Now().
func TimestampFunc(f func() time.Time) Option {
	return func(p *Parser) {
		p.timestampFunc = f
	}
}

// SubtestMode configures how Go subtests should be handled by the parser.
type SubtestMode string

const (
	// SubtestModeDefault is the default subtest mode. It treats tests with
	// subtests as any other tests.
	SubtestModeDefault SubtestMode = ""

	// IgnoreParentResults ignores test results for tests with subtests. Use
	// this mode if you use subtest parents for common setup/teardown, but are
	// not interested in counting them as failed tests. Ignoring their results
	// still preserves these tests and their captured output in the report.
	IgnoreParentResults SubtestMode = "ignore-parent-results"

	// ExcludeParents excludes tests that contain subtests from the report.
	// Note that the subtests themselves are not removed. Use this mode if you
	// use subtest parents for common setup/teardown, but are not actually
	// interested in their presence in the created report. If output was
	// captured for tests that are removed, the output is preserved in the
	// global report output.
	ExcludeParents SubtestMode = "exclude-parents"
)

// ParseSubtestMode returns a SubtestMode for the given string.
func ParseSubtestMode(in string) (SubtestMode, error) {
	switch in {
	case string(IgnoreParentResults):
		return IgnoreParentResults, nil
	case string(ExcludeParents):
		return ExcludeParents, nil
	default:
		return SubtestModeDefault, fmt.Errorf("unknown subtest mode: %v", in)
	}
}

// SetSubtestMode is an Option to change how the parser handles tests with
// subtests. See the documentation for the individual SubtestModes for more
// information.
func SetSubtestMode(mode SubtestMode) Option {
	return func(p *Parser) {
		p.subtestMode = mode
	}
}

// NewParser returns a new Go test output parser.
func NewParser(options ...Option) *Parser {
	p := &Parser{}
	for _, option := range options {
		option(p)
	}
	return p
}

// Parser is a Go test output Parser.
type Parser struct {
	packageName string
	subtestMode SubtestMode

	timestampFunc func() time.Time

	events []Event
}

// Parse parses Go test output from the given io.Reader r and returns
// gtr.Report.
func (p *Parser) Parse(r io.Reader) (gtr.Report, error) {
	p.events = nil
	s := bufio.NewScanner(r)
	for s.Scan() {
		p.parseLine(s.Text())
	}
	return p.report(p.events), s.Err()
}

// report generates a gtr.Report from the given list of events.
func (p *Parser) report(events []Event) gtr.Report {
	rb := newReportBuilder()
	rb.packageName = p.packageName
	rb.subtestMode = p.subtestMode
	if p.timestampFunc != nil {
		rb.timestampFunc = p.timestampFunc
	}
	for _, ev := range events {
		switch ev.Type {
		case "run_test":
			rb.CreateTest(ev.Name)
		case "pause_test":
			rb.PauseTest(ev.Name)
		case "cont_test":
			rb.ContinueTest(ev.Name)
		case "end_test":
			rb.EndTest(ev.Name, ev.Result, ev.Duration, ev.Indent)
		case "run_benchmark":
			rb.CreateBenchmark(ev.Name)
		case "benchmark":
			rb.BenchmarkResult(ev.Name, ev.Iterations, ev.NsPerOp, ev.MBPerSec, ev.BytesPerOp, ev.AllocsPerOp)
		case "end_benchmark":
			rb.EndBenchmark(ev.Name, ev.Result)
		case "status":
			rb.End()
		case "summary":
			rb.CreatePackage(ev.Name, ev.Result, ev.Duration, ev.Data)
		case "coverage":
			rb.Coverage(ev.CovPct, ev.CovPackages)
		case "build_output":
			rb.CreateBuildError(ev.Name)
		case "output":
			rb.AppendOutput(ev.Data)
		default:
			fmt.Printf("unhandled event type: %v\n", ev.Type)
		}
	}
	return rb.Build()
}

// Events returns the events created by the parser.
func (p *Parser) Events() []Event {
	events := make([]Event, len(p.events))
	copy(events, p.events)
	return events
}

func (p *Parser) parseLine(line string) {
	if strings.HasPrefix(line, "=== RUN ") {
		p.runTest(strings.TrimSpace(line[8:]))
	} else if strings.HasPrefix(line, "=== PAUSE ") {
		p.pauseTest(strings.TrimSpace(line[10:]))
	} else if strings.HasPrefix(line, "=== CONT ") {
		p.contTest(strings.TrimSpace(line[9:]))
	} else if matches := regexEndTest.FindStringSubmatch(line); len(matches) == 5 {
		p.endTest(line, matches[1], matches[2], matches[3], matches[4])
	} else if matches := regexStatus.FindStringSubmatch(line); len(matches) == 2 {
		p.status(matches[1])
	} else if matches := regexSummary.FindStringSubmatch(line); len(matches) == 8 {
		p.summary(matches[1], matches[2], matches[3], matches[4], matches[5], matches[6], matches[7])
	} else if matches := regexCoverage.FindStringSubmatch(line); len(matches) == 3 {
		p.coverage(matches[1], matches[2])
	} else if matches := regexBenchmark.FindStringSubmatch(line); len(matches) == 2 {
		p.runBench(matches[1])
	} else if matches := regexBenchSummary.FindStringSubmatch(line); len(matches) == 7 {
		p.benchSummary(matches[1], matches[2], matches[3], matches[4], matches[5], matches[6])
	} else if matches := regexEndBenchmark.FindStringSubmatch(line); len(matches) == 3 {
		p.endBench(matches[1], matches[2])
	} else if strings.HasPrefix(line, "# ") {
		// TODO(jstemmer): this should just be output; we should detect build output when building report
		fields := strings.Fields(strings.TrimPrefix(line, "# "))
		if len(fields) == 1 || len(fields) == 2 {
			p.buildOutput(fields[0])
		} else {
			p.output(line)
		}
	} else {
		p.output(line)
	}
}

func (p *Parser) add(event Event) {
	p.events = append(p.events, event)
}

func (p *Parser) runTest(name string) {
	p.add(Event{Type: "run_test", Name: name})
}

func (p *Parser) pauseTest(name string) {
	p.add(Event{Type: "pause_test", Name: name})
}

func (p *Parser) contTest(name string) {
	p.add(Event{Type: "cont_test", Name: name})
}

func (p *Parser) endTest(line, indent, result, name, duration string) {
	if idx := strings.Index(line, fmt.Sprintf("%s--- %s:", indent, result)); idx > 0 {
		p.output(line[:idx])
	}
	_, n := stripIndent(indent)
	p.add(Event{
		Type:     "end_test",
		Name:     name,
		Result:   result,
		Indent:   n,
		Duration: parseSeconds(duration),
	})
}

func (p *Parser) status(result string) {
	p.add(Event{Type: "status", Result: result})
}

func (p *Parser) summary(result, name, duration, cached, status, covpct, packages string) {
	p.add(Event{
		Type:        "summary",
		Result:      result,
		Name:        name,
		Duration:    parseSeconds(duration),
		Data:        strings.TrimSpace(cached + " " + status),
		CovPct:      parseFloat(covpct),
		CovPackages: parsePackages(packages),
	})
}

func (p *Parser) coverage(percent, packages string) {
	p.add(Event{
		Type:        "coverage",
		CovPct:      parseFloat(percent),
		CovPackages: parsePackages(packages),
	})
}

func (p *Parser) runBench(name string) {
	p.add(Event{
		Type: "run_benchmark",
		Name: name,
	})
}

func (p *Parser) benchSummary(name, iterations, nsPerOp, mbPerSec, bytesPerOp, allocsPerOp string) {
	p.add(Event{
		Type:        "benchmark",
		Name:        name,
		Iterations:  parseInt(iterations),
		NsPerOp:     parseFloat(nsPerOp),
		MBPerSec:    parseFloat(mbPerSec),
		BytesPerOp:  parseInt(bytesPerOp),
		AllocsPerOp: parseInt(allocsPerOp),
	})
}

func (p *Parser) endBench(result, name string) {
	p.add(Event{
		Type:   "end_benchmark",
		Name:   name,
		Result: result,
	})
}

func (p *Parser) buildOutput(packageName string) {
	p.add(Event{
		Type: "build_output",
		Name: packageName,
	})
}

func (p *Parser) output(line string) {
	p.add(Event{Type: "output", Data: line})
}

func parseSeconds(s string) time.Duration {
	if s == "" {
		return time.Duration(0)
	}
	// ignore error
	d, _ := time.ParseDuration(s + "s")
	return d
}

func parseFloat(s string) float64 {
	if s == "" {
		return 0
	}
	// ignore error
	pct, _ := strconv.ParseFloat(s, 64)
	return pct
}

func parsePackages(pkgList string) []string {
	if len(pkgList) == 0 {
		return nil
	}
	return strings.Split(pkgList, ", ")
}

func parseInt(s string) int64 {
	// ignore error
	n, _ := strconv.ParseInt(s, 10, 64)
	return n
}

func stripIndent(line string) (string, int) {
	var indent int
	for indent = 0; strings.HasPrefix(line, "    "); indent++ {
		line = line[4:]
	}
	return line, indent
}
<commit_msg>parser/gotest: Move NewParser function below Parser struct<commit_after>// Package gotest is a standard Go test output parser.
package gotest

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
)

var (
	// regexBenchInfo captures 3-5 groups: benchmark name, number of times ran, ns/op (with or without decimal), MB/sec (optional), B/op (optional), and allocs/op (optional).
	regexBenchmark    = regexp.MustCompile(`^(Benchmark[^ -]+)$`)
	regexBenchSummary = regexp.MustCompile(`^(Benchmark[^ -]+)(?:-\d+\s+|\s+)(\d+)\s+(\d+|\d+\.\d+)\sns\/op(?:\s+(\d+|\d+\.\d+)\sMB\/s)?(?:\s+(\d+)\sB\/op)?(?:\s+(\d+)\sallocs/op)?`)
	regexCoverage     = regexp.MustCompile(`^coverage:\s+(\d+|\d+\.\d+)%\s+of\s+statements(?:\sin\s(.+))?$`)
	regexEndBenchmark = regexp.MustCompile(`^--- (BENCH|FAIL|SKIP): (Benchmark[^ -]+)(?:-\d+)?$`)
	regexEndTest      = regexp.MustCompile(`((?:    )*)--- (PASS|FAIL|SKIP): ([^ ]+) \((\d+\.\d+)(?: seconds|s)\)`)
	regexStatus       = regexp.MustCompile(`^(PASS|FAIL|SKIP)$`)
	regexSummary      = regexp.MustCompile(`` +
		// 1: result
		`^(\?|ok|FAIL)` +
		// 2: package name
		`\s+([^ \t]+)` +
		// 3: duration (optional)
		`(?:\s+(\d+\.\d+)s)?` +
		// 4: cached indicator (optional)
		`(?:\s+(\(cached\)))?` +
		// 5: [status message] (optional)
		`(?:\s+(\[[^\]]+\]))?` +
		// 6: coverage percentage (optional)
		// 7: coverage package list (optional)
		`(?:\s+coverage:\s+(\d+\.\d+)%\sof\sstatements(?:\sin\s(.+))?)?$`)
)

// Option defines options that can be passed to gotest.New.
type Option func(*Parser)

// PackageName is an Option that sets the default package name to use when it
// cannot be determined from the test output.
func PackageName(name string) Option {
	return func(p *Parser) {
		p.packageName = name
	}
}

// TimestampFunc is an Option that sets the timestamp function that is used to
// determine the current time when creating the Report. This can be used to
// override the default behaviour of using time.Now().
func TimestampFunc(f func() time.Time) Option {
	return func(p *Parser) {
		p.timestampFunc = f
	}
}

// SubtestMode configures how Go subtests should be handled by the parser.
type SubtestMode string

const (
	// SubtestModeDefault is the default subtest mode. It treats tests with
	// subtests as any other tests.
	SubtestModeDefault SubtestMode = ""

	// IgnoreParentResults ignores test results for tests with subtests. Use
	// this mode if you use subtest parents for common setup/teardown, but are
	// not interested in counting them as failed tests. Ignoring their results
	// still preserves these tests and their captured output in the report.
	IgnoreParentResults SubtestMode = "ignore-parent-results"

	// ExcludeParents excludes tests that contain subtests from the report.
	// Note that the subtests themselves are not removed. Use this mode if you
	// use subtest parents for common setup/teardown, but are not actually
	// interested in their presence in the created report. If output was
	// captured for tests that are removed, the output is preserved in the
	// global report output.
	ExcludeParents SubtestMode = "exclude-parents"
)

// ParseSubtestMode returns a SubtestMode for the given string.
func ParseSubtestMode(in string) (SubtestMode, error) {
	switch in {
	case string(IgnoreParentResults):
		return IgnoreParentResults, nil
	case string(ExcludeParents):
		return ExcludeParents, nil
	default:
		return SubtestModeDefault, fmt.Errorf("unknown subtest mode: %v", in)
	}
}

// SetSubtestMode is an Option to change how the parser handles tests with
// subtests. See the documentation for the individual SubtestModes for more
// information.
func SetSubtestMode(mode SubtestMode) Option {
	return func(p *Parser) {
		p.subtestMode = mode
	}
}

// Parser is a Go test output Parser.
type Parser struct {
	packageName string
	subtestMode SubtestMode

	timestampFunc func() time.Time

	events []Event
}

// NewParser returns a new Go test output parser.
func NewParser(options ...Option) *Parser {
	p := &Parser{}
	for _, option := range options {
		option(p)
	}
	return p
}

// Parse parses Go test output from the given io.Reader r and returns
// gtr.Report.
func (p *Parser) Parse(r io.Reader) (gtr.Report, error) {
	p.events = nil
	s := bufio.NewScanner(r)
	for s.Scan() {
		p.parseLine(s.Text())
	}
	return p.report(p.events), s.Err()
}

// report generates a gtr.Report from the given list of events.
func (p *Parser) report(events []Event) gtr.Report {
	rb := newReportBuilder()
	rb.packageName = p.packageName
	rb.subtestMode = p.subtestMode
	if p.timestampFunc != nil {
		rb.timestampFunc = p.timestampFunc
	}
	for _, ev := range events {
		switch ev.Type {
		case "run_test":
			rb.CreateTest(ev.Name)
		case "pause_test":
			rb.PauseTest(ev.Name)
		case "cont_test":
			rb.ContinueTest(ev.Name)
		case "end_test":
			rb.EndTest(ev.Name, ev.Result, ev.Duration, ev.Indent)
		case "run_benchmark":
			rb.CreateBenchmark(ev.Name)
		case "benchmark":
			rb.BenchmarkResult(ev.Name, ev.Iterations, ev.NsPerOp, ev.MBPerSec, ev.BytesPerOp, ev.AllocsPerOp)
		case "end_benchmark":
			rb.EndBenchmark(ev.Name, ev.Result)
		case "status":
			rb.End()
		case "summary":
			rb.CreatePackage(ev.Name, ev.Result, ev.Duration, ev.Data)
		case "coverage":
			rb.Coverage(ev.CovPct, ev.CovPackages)
		case "build_output":
			rb.CreateBuildError(ev.Name)
		case "output":
			rb.AppendOutput(ev.Data)
		default:
			fmt.Printf("unhandled event type: %v\n", ev.Type)
		}
	}
	return rb.Build()
}

// Events returns the events created by the parser.
func (p *Parser) Events() []Event {
	events := make([]Event, len(p.events))
	copy(events, p.events)
	return events
}

func (p *Parser) parseLine(line string) {
	if strings.HasPrefix(line, "=== RUN ") {
		p.runTest(strings.TrimSpace(line[8:]))
	} else if strings.HasPrefix(line, "=== PAUSE ") {
		p.pauseTest(strings.TrimSpace(line[10:]))
	} else if strings.HasPrefix(line, "=== CONT ") {
		p.contTest(strings.TrimSpace(line[9:]))
	} else if matches := regexEndTest.FindStringSubmatch(line); len(matches) == 5 {
		p.endTest(line, matches[1], matches[2], matches[3], matches[4])
	} else if matches := regexStatus.FindStringSubmatch(line); len(matches) == 2 {
		p.status(matches[1])
	} else if matches := regexSummary.FindStringSubmatch(line); len(matches) == 8 {
		p.summary(matches[1], matches[2], matches[3], matches[4], matches[5], matches[6], matches[7])
	} else if matches := regexCoverage.FindStringSubmatch(line); len(matches) == 3 {
		p.coverage(matches[1], matches[2])
	} else if matches := regexBenchmark.FindStringSubmatch(line); len(matches) == 2 {
		p.runBench(matches[1])
	} else if matches := regexBenchSummary.FindStringSubmatch(line); len(matches) == 7 {
		p.benchSummary(matches[1], matches[2], matches[3], matches[4], matches[5], matches[6])
	} else if matches := regexEndBenchmark.FindStringSubmatch(line); len(matches) == 3 {
		p.endBench(matches[1], matches[2])
	} else if strings.HasPrefix(line, "# ") {
		// TODO(jstemmer): this should just be output; we should detect build output when building report
		fields := strings.Fields(strings.TrimPrefix(line, "# "))
		if len(fields) == 1 || len(fields) == 2 {
			p.buildOutput(fields[0])
		} else {
			p.output(line)
		}
	} else {
		p.output(line)
	}
}

func (p *Parser) add(event Event) {
	p.events = append(p.events, event)
}

func (p *Parser) runTest(name string) {
	p.add(Event{Type: "run_test", Name: name})
}

func (p *Parser) pauseTest(name string) {
	p.add(Event{Type: "pause_test", Name: name})
}

func (p *Parser) contTest(name string) {
	p.add(Event{Type: "cont_test", Name: name})
}

func (p *Parser) endTest(line, indent, result, name, duration string) {
	if idx := strings.Index(line, fmt.Sprintf("%s--- %s:", indent, result)); idx > 0 {
		p.output(line[:idx])
	}
	_, n := stripIndent(indent)
	p.add(Event{
		Type:     "end_test",
		Name:     name,
		Result:   result,
		Indent:   n,
		Duration: parseSeconds(duration),
	})
}

func (p *Parser) status(result string) {
	p.add(Event{Type: "status", Result: result})
}

func (p *Parser) summary(result, name, duration, cached, status, covpct, packages string) {
	p.add(Event{
		Type:        "summary",
		Result:      result,
		Name:        name,
		Duration:    parseSeconds(duration),
		Data:        strings.TrimSpace(cached + " " + status),
		CovPct:      parseFloat(covpct),
		CovPackages: parsePackages(packages),
	})
}

func (p *Parser) coverage(percent, packages string) {
	p.add(Event{
		Type:        "coverage",
		CovPct:      parseFloat(percent),
		CovPackages: parsePackages(packages),
	})
}

func (p *Parser) runBench(name string) {
	p.add(Event{
		Type: "run_benchmark",
		Name: name,
	})
}

func (p *Parser) benchSummary(name, iterations, nsPerOp, mbPerSec, bytesPerOp, allocsPerOp string) {
	p.add(Event{
		Type:        "benchmark",
		Name:        name,
		Iterations:  parseInt(iterations),
		NsPerOp:     parseFloat(nsPerOp),
		MBPerSec:    parseFloat(mbPerSec),
		BytesPerOp:  parseInt(bytesPerOp),
		AllocsPerOp: parseInt(allocsPerOp),
	})
}

func (p *Parser) endBench(result, name string) {
	p.add(Event{
		Type:   "end_benchmark",
		Name:   name,
		Result: result,
	})
}

func (p *Parser) buildOutput(packageName string) {
	p.add(Event{
		Type: "build_output",
		Name: packageName,
	})
}

func (p *Parser) output(line string) {
	p.add(Event{Type: "output", Data: line})
}

func parseSeconds(s string) time.Duration {
	if s == "" {
		return time.Duration(0)
	}
	// ignore error
	d, _ := time.ParseDuration(s + "s")
	return d
}

func parseFloat(s string) float64 {
	if s == "" {
		return 0
	}
	// ignore error
	pct, _ := strconv.ParseFloat(s, 64)
	return pct
}

func parsePackages(pkgList string) []string {
	if len(pkgList) == 0 {
		return nil
	}
	return strings.Split(pkgList, ", ")
}

func parseInt(s string) int64 {
	// ignore error
	n, _ := strconv.ParseInt(s, 10, 64)
	return n
}

func stripIndent(line string) (string, int) {
	var indent int
	for indent = 0; strings.HasPrefix(line, "    "); indent++ {
		line = line[4:]
	}
	return line, indent
}
<|endoftext|>


<commit_before>package main

import (
	"log"

	"github.com/byuoitav/hateoas"
	"github.com/byuoitav/sony-control-microservice/handlers"
	"github.com/byuoitav/wso2jwt"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/engine/fasthttp"
	"github.com/labstack/echo/middleware"
)

func main() {
	err := hateoas.Load("https://raw.githubusercontent.com/byuoitav/sony-control-microservice/master/swagger.json")
	if err != nil {
		log.Fatalln("Could not load swagger.json file. Error: " + err.Error())
	}

	port := ":8007"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())

	router.Get("/", hateoas.RootResponse)
	router.Get("/health", health.Check)

	router.Get("/:address/list/commands", handlers.GetCommands, wso2jwt.ValidateJWT())
	router.Get("/:address/command/:command", handlers.SendCommand, wso2jwt.ValidateJWT())
	router.Get("/:address/command/:command/count/:count", handlers.FloodCommand, wso2jwt.ValidateJWT())

	router.Get("/:address/power/on", handlers.PowerOn, wso2jwt.ValidateJWT())
	router.Get("/:address/power/standby", handlers.Standby, wso2jwt.ValidateJWT())
	router.Get("/:address/input/:port", handlers.SwitchInput, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/set/:difference", handlers.SetVolume, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/calibrate/:default", handlers.CalibrateVolume, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/up", handlers.VolumeUp, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/down", handlers.VolumeDown, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/mute", handlers.VolumeMute, wso2jwt.ValidateJWT())
	router.Get("/:address/volume/unmute", handlers.VolumeUnmute, wso2jwt.ValidateJWT())
	router.Get("/:address/display/blank", handlers.BlankDisplay, wso2jwt.ValidateJWT())
	router.Get("/:address/display/unblank", handlers.UnblankDisplay, wso2jwt.ValidateJWT())

	log.Println("Sony Control microservice is listening on " + port)
	server := fasthttp.New(port)
	server.ReadBufferSize = 1024 * 10 // Needed to interface properly with WSO2
	router.Run(server)
}
<commit_msg>Dork boy<commit_after>package main

import (
	"log"
	"net/http"

	"github.com/byuoitav/authmiddleware"
	"github.com/byuoitav/hateoas"
	"github.com/byuoitav/sony-control-microservice/handlers"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

func main() {
	err := hateoas.Load("https://raw.githubusercontent.com/byuoitav/sony-control-microservice/master/swagger.json")
	if err != nil {
		log.Fatalln("Could not load swagger.json file. Error: " + err.Error())
	}

	port := ":8007"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())

	// Use the `secure` routing group to require authentication
	secure := router.Group("", echo.WrapMiddleware(authmiddleware.Authenticate))

	router.GET("/", echo.WrapHandler(http.HandlerFunc(hateoas.RootResponse)))
	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))

	secure.GET("/:address/list/commands", handlers.GetCommands)
	secure.GET("/:address/command/:command", handlers.SendCommand)
	secure.GET("/:address/command/:command/count/:count", handlers.FloodCommand)

	secure.GET("/:address/power/on", handlers.PowerOn)
	secure.GET("/:address/power/standby", handlers.Standby)
	secure.GET("/:address/input/:port", handlers.SwitchInput)
	secure.GET("/:address/volume/set/:difference", handlers.SetVolume)
	secure.GET("/:address/volume/calibrate/:default", handlers.CalibrateVolume)
	secure.GET("/:address/volume/up", handlers.VolumeUp)
	secure.GET("/:address/volume/down", handlers.VolumeDown)
	secure.GET("/:address/volume/mute", handlers.VolumeMute)
	secure.GET("/:address/volume/unmute", handlers.VolumeUnmute)
	secure.GET("/:address/display/blank", handlers.BlankDisplay)
	secure.GET("/:address/display/unblank", handlers.UnblankDisplay)

	server := http.Server{
		Addr:           port,
		MaxHeaderBytes: 1024 * 10,
	}

	router.StartServer(&server)
}
<|endoftext|>