<commit_before>package server

import (
	"os"
	"path"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/coreos/etcd/clientv3"
	storagepb "github.com/coreos/etcd/storage/storagepb"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/pdpb"
)

// IsLeader returns whether server is leader or not.
func (s *Server) IsLeader() bool {
	return atomic.LoadInt64(&s.isLeader) == 1
}

func (s *Server) enableLeader(b bool) {
	value := int64(0)
	if b {
		value = 1
	}

	atomic.StoreInt64(&s.isLeader, value)

	if !b {
		// if we lost leader, we may:
		//	1, close all client connections
		//	2, close all running raft clusters
		s.closeAllConnections()

		s.closeClusters()
	}
}

// GetLeaderPath returns the leader path.
func GetLeaderPath(rootPath string) string {
	return path.Join(rootPath, "leader")
}

func (s *Server) getLeaderPath() string {
	return GetLeaderPath(s.cfg.RootPath)
}

func (s *Server) leaderLoop() {
	defer func() {
		s.wg.Done()
		s.enableLeader(false)
	}()

	for {
		if s.IsClosed() {
			return
		}

		s.enableLeader(false)

		leader, err := s.getLeader()
		if err != nil {
			log.Errorf("get leader err %v", err)
			continue
		} else if leader != nil {
			log.Debugf("leader is %s, watch it", leader.String())

			s.watchLeader()

			log.Debugf("leader changed, try to campaign leader")
		}

		if err = s.campaignLeader(); err != nil {
			log.Errorf("campaign leader err %s", err)
		}
	}
}

// GetLeader gets server leader from etcd.
func GetLeader(c *clientv3.Client, leaderPath string) (*pdpb.Leader, error) {
	leader := pdpb.Leader{}
	ok, err := getProtoMsg(c, leaderPath, &leader)
	if err != nil || !ok {
		return nil, errors.Trace(err)
	}

	return &leader, nil
}

func (s *Server) getLeader() (*pdpb.Leader, error) {
	return GetLeader(s.client, s.getLeaderPath())
}

func (s *Server) marshalLeader() string {
	leader := &pdpb.Leader{
		Addr: proto.String(s.cfg.Addr),
		Pid:  proto.Int64(int64(os.Getpid())),
	}

	data, err := proto.Marshal(leader)
	if err != nil {
		// can't fail, so panic here.
		log.Fatalf("marshal leader %v err %v", leader, err)
	}

	return string(data)
}

func (s *Server) campaignLeader() error {
	log.Debug("begin to campaign leader")

	lessor := clientv3.NewLease(s.client)
	defer lessor.Close()

	leaseResp, err := lessor.Create(context.TODO(), s.cfg.LeaderLease)
	if err != nil {
		return errors.Trace(err)
	}

	leaderKey := s.getLeaderPath()
	// The leader key must not exist, so the CreatedRevision is 0.
	resp, err := s.client.Txn(context.TODO()).
		If(clientv3.Compare(clientv3.CreatedRevision(leaderKey), "=", 0)).
		Then(clientv3.OpPut(leaderKey, s.leaderValue, clientv3.WithLease(clientv3.LeaseID(leaseResp.ID)))).
		Commit()
	if err != nil {
		return errors.Trace(err)
	} else if !resp.Succeeded {
		return errors.New("campaign leader failed, other server may campaign ok")
	}

	log.Debug("campaign leader ok")
	s.enableLeader(true)
	defer s.enableLeader(false)

	// keeps the leader
	ch, err := lessor.KeepAlive(s.client.Ctx(), clientv3.LeaseID(leaseResp.ID))
	if err != nil {
		return errors.Trace(err)
	}

	log.Debug("sync timestamp for tso")
	if err = s.syncTimestamp(); err != nil {
		return errors.Trace(err)
	}

	tsTicker := time.NewTicker(time.Duration(updateTimestampStep) * time.Millisecond)
	defer tsTicker.Stop()

	for {
		select {
		case _, ok := <-ch:
			if !ok {
				log.Infof("keep alive channel is closed")
				return nil
			}
		case <-tsTicker.C:
			if err = s.updateTimestamp(); err != nil {
				return errors.Trace(err)
			}
		}
	}
}

func (s *Server) watchLeader() {
	watcher := clientv3.NewWatcher(s.client)
	defer watcher.Close()

	for {
		rch := watcher.Watch(s.client.Ctx(), s.getLeaderPath())
		for wresp := range rch {
			if wresp.Canceled {
				return
			}

			for _, ev := range wresp.Events {
				if ev.Type == storagepb.EXPIRE || ev.Type == storagepb.DELETE {
					log.Infof("leader is expired or deleted")
					return
				}
			}
		}
	}
}

func (s *Server) leaderCmp() clientv3.Cmp {
	return clientv3.Compare(clientv3.Value(s.getLeaderPath()), "=", s.leaderValue)
}
<commit_msg>server: Address comment.<commit_after>package server

import (
	"os"
	"path"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/coreos/etcd/clientv3"
	storagepb "github.com/coreos/etcd/storage/storagepb"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/pdpb"
)

// IsLeader returns whether server is leader or not.
func (s *Server) IsLeader() bool {
	return atomic.LoadInt64(&s.isLeader) == 1
}

func (s *Server) enableLeader(b bool) {
	value := int64(0)
	if b {
		value = 1
	}

	atomic.StoreInt64(&s.isLeader, value)

	if !b {
		// if we lost leader, we may:
		//	1, close all client connections
		//	2, close all running raft clusters
		s.closeAllConnections()

		s.closeClusters()
	}
}

// GetLeaderPath returns the leader path.
func GetLeaderPath(rootPath string) string {
	return path.Join(rootPath, "leader")
}

func (s *Server) getLeaderPath() string {
	return GetLeaderPath(s.cfg.RootPath)
}

func (s *Server) leaderLoop() {
	defer func() {
		s.wg.Done()
		s.enableLeader(false)
	}()

	for {
		if s.IsClosed() {
			return
		}

		s.enableLeader(false)

		leader, err := s.getLeader()
		if err != nil {
			log.Errorf("get leader err %v", err)
			continue
		} else if leader != nil {
			log.Debugf("leader is %s, watch it", leader)

			s.watchLeader()

			log.Debugf("leader changed, try to campaign leader")
		}

		if err = s.campaignLeader(); err != nil {
			log.Errorf("campaign leader err %s", err)
		}
	}
}

// GetLeader gets server leader from etcd.
func GetLeader(c *clientv3.Client, leaderPath string) (*pdpb.Leader, error) {
	leader := pdpb.Leader{}
	ok, err := getProtoMsg(c, leaderPath, &leader)
	if err != nil || !ok {
		return nil, errors.Trace(err)
	}

	return &leader, nil
}

func (s *Server) getLeader() (*pdpb.Leader, error) {
	return GetLeader(s.client, s.getLeaderPath())
}

func (s *Server) marshalLeader() string {
	leader := &pdpb.Leader{
		Addr: proto.String(s.cfg.Addr),
		Pid:  proto.Int64(int64(os.Getpid())),
	}

	data, err := proto.Marshal(leader)
	if err != nil {
		// can't fail, so panic here.
		log.Fatalf("marshal leader %s err %v", leader, err)
	}

	return string(data)
}

func (s *Server) campaignLeader() error {
	log.Debug("begin to campaign leader")

	lessor := clientv3.NewLease(s.client)
	defer lessor.Close()

	leaseResp, err := lessor.Create(context.TODO(), s.cfg.LeaderLease)
	if err != nil {
		return errors.Trace(err)
	}

	leaderKey := s.getLeaderPath()
	// The leader key must not exist, so the CreatedRevision is 0.
	resp, err := s.client.Txn(context.TODO()).
		If(clientv3.Compare(clientv3.CreatedRevision(leaderKey), "=", 0)).
		Then(clientv3.OpPut(leaderKey, s.leaderValue, clientv3.WithLease(clientv3.LeaseID(leaseResp.ID)))).
		Commit()
	if err != nil {
		return errors.Trace(err)
	} else if !resp.Succeeded {
		return errors.New("campaign leader failed, other server may campaign ok")
	}

	log.Debug("campaign leader ok")
	s.enableLeader(true)
	defer s.enableLeader(false)

	// keeps the leader
	ch, err := lessor.KeepAlive(s.client.Ctx(), clientv3.LeaseID(leaseResp.ID))
	if err != nil {
		return errors.Trace(err)
	}

	log.Debug("sync timestamp for tso")
	if err = s.syncTimestamp(); err != nil {
		return errors.Trace(err)
	}

	tsTicker := time.NewTicker(time.Duration(updateTimestampStep) * time.Millisecond)
	defer tsTicker.Stop()

	for {
		select {
		case _, ok := <-ch:
			if !ok {
				log.Infof("keep alive channel is closed")
				return nil
			}
		case <-tsTicker.C:
			if err = s.updateTimestamp(); err != nil {
				return errors.Trace(err)
			}
		}
	}
}

func (s *Server) watchLeader() {
	watcher := clientv3.NewWatcher(s.client)
	defer watcher.Close()

	for {
		rch := watcher.Watch(s.client.Ctx(), s.getLeaderPath())
		for wresp := range rch {
			if wresp.Canceled {
				return
			}

			for _, ev := range wresp.Events {
				if ev.Type == storagepb.EXPIRE || ev.Type == storagepb.DELETE {
					log.Infof("leader is expired or deleted")
					return
				}
			}
		}
	}
}

func (s *Server) leaderCmp() clientv3.Cmp {
	return clientv3.Compare(clientv3.Value(s.getLeaderPath()), "=", s.leaderValue)
}
<|endoftext|>


<commit_before>package server

import (
	"os"
	"path"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/lease"
	storagepb "github.com/coreos/etcd/storage/storagepb"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/pdpb"
)

// IsLeader returns whether server is leader or not.
func (s *Server) IsLeader() bool {
	return atomic.LoadInt64(&s.isLeader) == 1
}

func (s *Server) enableLeader(b bool) {
	value := int64(0)
	if b {
		value = 1
	}

	atomic.StoreInt64(&s.isLeader, value)

	if !b {
		// if we lost leader, we may:
		//	1, close all client connections
		//	2, close all running raft clusters
		s.closeAllConnections()

		s.closeClusters()
	}
}

// GetLeaderPath returns the leader path.
func GetLeaderPath(rootPath string) string {
	return path.Join(rootPath, "leader")
}

func (s *Server) getLeaderPath() string {
	return GetLeaderPath(s.cfg.RootPath)
}

func (s *Server) leaderLoop() {
	defer func() {
		s.wg.Done()
		s.enableLeader(false)
	}()

	for {
		if s.IsClosed() {
			return
		}

		s.enableLeader(false)

		leader, err := s.getLeader()
		if err != nil {
			log.Errorf("get leader err %v", err)
			continue
		} else if leader != nil {
			log.Debugf("leader is %#v, watch it", leader)

			s.watchLeader()

			log.Debugf("leader changed, try to campaign leader")
		}

		if err = s.campaignLeader(); err != nil {
			log.Errorf("campaign leader err %s", err)
		}
	}
}

// GetLeader gets server leader from etcd.
func GetLeader(c *clientv3.Client, leaderPath string) (*pdpb.Leader, error) {
	leader := pdpb.Leader{}
	ok, err := getProtoMsg(c, leaderPath, &leader)
	if err != nil || !ok {
		return nil, errors.Trace(err)
	}

	return &leader, nil
}

func (s *Server) getLeader() (*pdpb.Leader, error) {
	return GetLeader(s.client, s.getLeaderPath())
}

func (s *Server) marshalLeader() string {
	leader := &pdpb.Leader{
		Addr: proto.String(s.cfg.Addr),
		Pid:  proto.Int64(int64(os.Getpid())),
	}

	data, err := proto.Marshal(leader)
	if err != nil {
		// can't fail, so panic here.
		log.Fatalf("marshal leader %v err %v", leader, err)
	}

	return string(data)
}

func (s *Server) campaignLeader() error {
	log.Debug("begin to campaign leader")

	lessor := clientv3.NewLease(s.client)
	defer lessor.Close()

	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	leaseResp, err := lessor.Create(ctx, s.cfg.LeaderLease)
	if err != nil {
		return errors.Trace(err)
	}
	cancel()

	leaderKey := s.getLeaderPath()
	// The leader key must not exist, so the CreatedRevision is 0.
	ctx, cancel = context.WithTimeout(context.Background(), requestTimeout)
	resp, err := s.client.Txn(ctx).
		If(clientv3.Compare(clientv3.CreatedRevision(leaderKey), "=", 0)).
		Then(clientv3.OpPut(leaderKey, s.leaderValue, clientv3.WithLease(lease.LeaseID(leaseResp.ID)))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	} else if !resp.Succeeded {
		return errors.New("campaign leader failed, other server may campaign ok")
	}

	log.Debug("campaign leader ok")
	s.enableLeader(true)
	defer s.enableLeader(false)

	// keeps the leader
	ch, err := lessor.KeepAlive(s.client.Ctx(), lease.LeaseID(leaseResp.ID))
	if err != nil {
		return errors.Trace(err)
	}

	log.Debug("sync timestamp for tso")
	if err = s.syncTimestamp(); err != nil {
		return errors.Trace(err)
	}

	tsTicker := time.NewTicker(time.Duration(updateTimestampStep) * time.Millisecond)
	defer tsTicker.Stop()

	for {
		select {
		case _, ok := <-ch:
			if !ok {
				log.Infof("keep alive channel is closed")
				return nil
			}
		case <-tsTicker.C:
			if err = s.updateTimestamp(); err != nil {
				return errors.Trace(err)
			}
		}
	}
}

func (s *Server) watchLeader() {
	watcher := clientv3.NewWatcher(s.client)
	defer watcher.Close()

	for {
		rch := watcher.Watch(s.client.Ctx(), s.getLeaderPath())
		for wresp := range rch {
			if wresp.Canceled {
				return
			}

			for _, ev := range wresp.Events {
				if ev.Type == storagepb.EXPIRE || ev.Type == storagepb.DELETE {
					log.Infof("leader is expired or deleted")
					return
				}
			}
		}
	}
}

func (s *Server) leaderCmp() clientv3.Cmp {
	return clientv3.Compare(clientv3.Value(s.getLeaderPath()), "=", s.leaderValue)
}
<commit_msg>server: Address comment.<commit_after>package server

import (
	"os"
	"path"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/lease"
	storagepb "github.com/coreos/etcd/storage/storagepb"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/pdpb"
)

// IsLeader returns whether server is leader or not.
func (s *Server) IsLeader() bool {
	return atomic.LoadInt64(&s.isLeader) == 1
}

func (s *Server) enableLeader(b bool) {
	value := int64(0)
	if b {
		value = 1
	}

	atomic.StoreInt64(&s.isLeader, value)

	if !b {
		// if we lost leader, we may:
		//	1, close all client connections
		//	2, close all running raft clusters
		s.closeAllConnections()

		s.closeClusters()
	}
}

// GetLeaderPath returns the leader path.
func GetLeaderPath(rootPath string) string {
	return path.Join(rootPath, "leader")
}

func (s *Server) getLeaderPath() string {
	return GetLeaderPath(s.cfg.RootPath)
}

func (s *Server) leaderLoop() {
	defer func() {
		s.wg.Done()
		s.enableLeader(false)
	}()

	for {
		if s.IsClosed() {
			return
		}

		s.enableLeader(false)

		leader, err := s.getLeader()
		if err != nil {
			log.Errorf("get leader err %v", err)
			continue
		} else if leader != nil {
			log.Debugf("leader is %#v, watch it", leader)

			s.watchLeader()

			log.Debugf("leader changed, try to campaign leader")
		}

		if err = s.campaignLeader(); err != nil {
			log.Errorf("campaign leader err %s", err)
		}
	}
}

// GetLeader gets server leader from etcd.
func GetLeader(c *clientv3.Client, leaderPath string) (*pdpb.Leader, error) {
	leader := pdpb.Leader{}
	ok, err := getProtoMsg(c, leaderPath, &leader)
	if err != nil || !ok {
		return nil, errors.Trace(err)
	}

	return &leader, nil
}

func (s *Server) getLeader() (*pdpb.Leader, error) {
	return GetLeader(s.client, s.getLeaderPath())
}

func (s *Server) marshalLeader() string {
	leader := &pdpb.Leader{
		Addr: proto.String(s.cfg.Addr),
		Pid:  proto.Int64(int64(os.Getpid())),
	}

	data, err := proto.Marshal(leader)
	if err != nil {
		// can't fail, so panic here.
		log.Fatalf("marshal leader %v err %v", leader, err)
	}

	return string(data)
}

func (s *Server) campaignLeader() error {
	log.Debug("begin to campaign leader")

	lessor := clientv3.NewLease(s.client)
	defer lessor.Close()

	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	leaseResp, err := lessor.Create(ctx, s.cfg.LeaderLease)
	cancel()
	if err != nil {
		return errors.Trace(err)
	}

	leaderKey := s.getLeaderPath()
	// The leader key must not exist, so the CreatedRevision is 0.
	ctx, cancel = context.WithTimeout(context.Background(), requestTimeout)
	resp, err := s.client.Txn(ctx).
		If(clientv3.Compare(clientv3.CreatedRevision(leaderKey), "=", 0)).
		Then(clientv3.OpPut(leaderKey, s.leaderValue, clientv3.WithLease(lease.LeaseID(leaseResp.ID)))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	} else if !resp.Succeeded {
		return errors.New("campaign leader failed, other server may campaign ok")
	}

	log.Debug("campaign leader ok")
	s.enableLeader(true)
	defer s.enableLeader(false)

	// keeps the leader
	ch, err := lessor.KeepAlive(s.client.Ctx(), lease.LeaseID(leaseResp.ID))
	if err != nil {
		return errors.Trace(err)
	}

	log.Debug("sync timestamp for tso")
	if err = s.syncTimestamp(); err != nil {
		return errors.Trace(err)
	}

	tsTicker := time.NewTicker(time.Duration(updateTimestampStep) * time.Millisecond)
	defer tsTicker.Stop()

	for {
		select {
		case _, ok := <-ch:
			if !ok {
				log.Infof("keep alive channel is closed")
				return nil
			}
		case <-tsTicker.C:
			if err = s.updateTimestamp(); err != nil {
				return errors.Trace(err)
			}
		}
	}
}

func (s *Server) watchLeader() {
	watcher := clientv3.NewWatcher(s.client)
	defer watcher.Close()

	for {
		rch := watcher.Watch(s.client.Ctx(), s.getLeaderPath())
		for wresp := range rch {
			if wresp.Canceled {
				return
			}

			for _, ev := range wresp.Events {
				if ev.Type == storagepb.EXPIRE || ev.Type == storagepb.DELETE {
					log.Infof("leader is expired or deleted")
					return
				}
			}
		}
	}
}

func (s *Server) leaderCmp() clientv3.Cmp {
	return clientv3.Compare(clientv3.Value(s.getLeaderPath()), "=", s.leaderValue)
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"kmud/database"
	"kmud/game"
	"kmud/model"
	"kmud/utils"
	"labix.org/v2/mgo"
	"net"
	"strconv"
)

func login(conn net.Conn) *database.User {
	for {
		line := utils.GetUserInput(conn, "Username: ")

		if line == "" {
			return nil
		}

		user := model.M.GetUserByName(line)

		if user == nil {
			utils.WriteLine(conn, "User not found")
		} else if user.Online() {
			utils.WriteLine(conn, "That user is already online")
		} else {
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newUser(conn net.Conn) *database.User {
	for {
		name := utils.GetUserInput(conn, "Desired username: ")

		if name == "" {
			return nil
		}

		user := model.M.GetUserByName(name)

		if user != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			user = model.M.CreateUser(name)
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newCharacter(conn net.Conn, user *database.User) *database.Character {
	// TODO: character slot limit
	const SizeLimit = 12
	for {
		name := utils.GetUserInput(conn, "Desired character name: ")

		if name == "" {
			return nil
		}

		char := model.M.GetCharacterByName(name)

		if char != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			room := model.M.GetRooms()[0] // TODO: Better way to pick an initial character location
			return model.M.CreateCharacter(name, user, room)
		}
	}

	panic("Unexpected code path")
	return nil
}

func quit(conn net.Conn) {
	utils.WriteLine(conn, "Take luck!")
	conn.Close()
}

func mainMenu() utils.Menu {
	menu := utils.NewMenu("MUD")

	menu.AddAction("l", "[L]ogin")
	menu.AddAction("n", "[N]ew user")
	menu.AddAction("q", "[Q]uit")

	return menu
}

func userMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user)

	menu := utils.NewMenu(user.PrettyName())
	menu.AddAction("l", "[L]ogout")
	menu.AddAction("a", "[A]dmin")
	menu.AddAction("n", "[N]ew character")
	if len(chars) > 0 {
		menu.AddAction("d", "[D]elete character")
	}

	// TODO: Sort character list

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.GetId())
	}

	return menu
}

func deleteMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user)

	menu := utils.NewMenu("Delete character")

	menu.AddAction("c", "[C]ancel")

	// TODO: Sort character list

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.GetId())
	}

	return menu
}

func adminMenu() utils.Menu {
	menu := utils.NewMenu("Admin")
	menu.AddAction("u", "[U]sers")
	return menu
}

func userAdminMenu() utils.Menu {
	menu := utils.NewMenu("User Admin")

	for i, user := range model.M.GetUsers() {
		index := i + 1

		online := ""
		if user.Online() {
			online = "*"
		}

		actionText := fmt.Sprintf("[%v]%v", index, user.PrettyName()+online)
		menu.AddActionData(index, actionText, user.GetId())
	}

	return menu
}

func userSpecificMenu(user *database.User) utils.Menu {
	menu := utils.NewMenu("User: " + user.PrettyName())
	menu.AddAction("d", "[D]elete")
	return menu
}

func handleConnection(session *mgo.Session, conn net.Conn) {
	defer conn.Close()
	defer session.Close()

	var user *database.User
	var character *database.Character

	defer func() {
		if r := recover(); r != nil {
			username := ""
			charname := ""

			if user != nil {
				user.SetOnline(false)
				username = user.PrettyName()
			}

			if character != nil {
				charname = character.PrettyName()
			}

			fmt.Printf("Lost connection to client (%v/%v): %v, %v\n",
				username,
				charname,
				conn.RemoteAddr(),
				r)
		}
	}()

	for {
		if user == nil {
			menu := mainMenu()
			choice, _ := menu.Exec(conn, utils.ColorModeNone)

			switch choice {
			case "l":
				user = login(conn)
			case "n":
				user = newUser(conn)
			case "":
				fallthrough
			case "q":
				quit(conn)
				return
			}
		} else if character == nil {
			menu := userMenu(user)
			choice, charId := menu.Exec(conn, user.GetColorMode())

			switch choice {
			case "":
				fallthrough
			case "l":
				user.SetOnline(false)
				user = nil
			case "a":
				adminMenu := adminMenu()
				for {
					choice, _ := adminMenu.Exec(conn, user.GetColorMode())
					if choice == "" {
						break
					} else if choice == "u" {
						for {
							userAdminMenu := userAdminMenu()
							choice, userId := userAdminMenu.Exec(conn, user.GetColorMode())
							if choice == "" {
								break
							} else {
								_, err := strconv.Atoi(choice)

								if err == nil {
									for {
										userMenu := userSpecificMenu(model.M.GetUser(userId))
										choice, _ = userMenu.Exec(conn, user.GetColorMode())
										if choice == "" {
											break
										} else if choice == "d" {
											model.M.DeleteUser(userId)
											break
										}
									}
								}
							}
						}
					}
				}
			case "n":
				character = newCharacter(conn, user)
			case "d":
				for {
					deleteMenu := deleteMenu(user)
					deleteChoice, deleteCharId := deleteMenu.Exec(conn, user.GetColorMode())

					if deleteChoice == "" || deleteChoice == "c" {
						break
					}

					_, err := strconv.Atoi(deleteChoice)

					if err == nil {
						// TODO: Delete confirmation
						model.M.DeleteCharacter(deleteCharId)
					}
				}

			default:
				_, err := strconv.Atoi(choice)

				if err == nil {
					character = model.M.GetCharacter(charId)
				}
			}
		} else {
			game.Exec(conn, user, character)
			character = nil
		}
	}
}

func Exec() {
	fmt.Printf("Connecting to database... ")
	session, err := mgo.Dial("localhost")

	utils.HandleError(err)

	fmt.Println("done.")

	listener, err := net.Listen("tcp", ":8945")
	utils.HandleError(err)

	err = model.Init(session.Copy())

	// If there are no rooms at all create one
	rooms := model.M.GetRooms()
	if len(rooms) == 0 {
		zones := model.M.GetZones()

		var zone *database.Zone

		if len(zones) == 0 {
			zone = model.M.CreateZone("Default")
		} else {
			zone = zones[0]
		}

		model.M.CreateRoom(zone)
	}

	fmt.Println("Server listening on port 8945")

	for {
		conn, err := listener.Accept()
		utils.HandleError(err)
		fmt.Println("Client connected:", conn.RemoteAddr())
		go handleConnection(session.Copy(), conn)
	}
}

// vim: nocindent
<commit_msg>handleConnection() doesn't need a mgo.Session pointer anymore<commit_after>package server

import (
	"fmt"
	"kmud/database"
	"kmud/game"
	"kmud/model"
	"kmud/utils"
	"labix.org/v2/mgo"
	"net"
	"strconv"
)

func login(conn net.Conn) *database.User {
	for {
		line := utils.GetUserInput(conn, "Username: ")

		if line == "" {
			return nil
		}

		user := model.M.GetUserByName(line)

		if user == nil {
			utils.WriteLine(conn, "User not found")
		} else if user.Online() {
			utils.WriteLine(conn, "That user is already online")
		} else {
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newUser(conn net.Conn) *database.User {
	for {
		name := utils.GetUserInput(conn, "Desired username: ")

		if name == "" {
			return nil
		}

		user := model.M.GetUserByName(name)

		if user != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			user = model.M.CreateUser(name)
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newCharacter(conn net.Conn, user *database.User) *database.Character {
	// TODO: character slot limit
	const SizeLimit = 12
	for {
		name := utils.GetUserInput(conn, "Desired character name: ")

		if name == "" {
			return nil
		}

		char := model.M.GetCharacterByName(name)

		if char != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			room := model.M.GetRooms()[0] // TODO: Better way to pick an initial character location
			return model.M.CreateCharacter(name, user, room)
		}
	}

	panic("Unexpected code path")
	return nil
}

func quit(conn net.Conn) {
	utils.WriteLine(conn, "Take luck!")
	conn.Close()
}

func mainMenu() utils.Menu {
	menu := utils.NewMenu("MUD")

	menu.AddAction("l", "[L]ogin")
	menu.AddAction("n", "[N]ew user")
	menu.AddAction("q", "[Q]uit")

	return menu
}

func userMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user)

	menu := utils.NewMenu(user.PrettyName())
	menu.AddAction("l", "[L]ogout")
	menu.AddAction("a", "[A]dmin")
	menu.AddAction("n", "[N]ew character")
	if len(chars) > 0 {
		menu.AddAction("d", "[D]elete character")
	}

	// TODO: Sort character list

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.GetId())
	}

	return menu
}

func deleteMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user)

	menu := utils.NewMenu("Delete character")

	menu.AddAction("c", "[C]ancel")

	// TODO: Sort character list

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.GetId())
	}

	return menu
}

func adminMenu() utils.Menu {
	menu := utils.NewMenu("Admin")
	menu.AddAction("u", "[U]sers")
	return menu
}

func userAdminMenu() utils.Menu {
	menu := utils.NewMenu("User Admin")

	for i, user := range model.M.GetUsers() {
		index := i + 1

		online := ""
		if user.Online() {
			online = "*"
		}

		actionText := fmt.Sprintf("[%v]%v", index, user.PrettyName()+online)
		menu.AddActionData(index, actionText, user.GetId())
	}

	return menu
}

func userSpecificMenu(user *database.User) utils.Menu {
	menu := utils.NewMenu("User: " + user.PrettyName())
	menu.AddAction("d", "[D]elete")
	return menu
}

func handleConnection(conn net.Conn) {
	defer conn.Close()

	var user *database.User
	var character *database.Character

	defer func() {
		if r := recover(); r != nil {
			username := ""
			charname := ""

			if user != nil {
				user.SetOnline(false)
				username = user.PrettyName()
			}

			if character != nil {
				charname = character.PrettyName()
			}

			fmt.Printf("Lost connection to client (%v/%v): %v, %v\n",
				username,
				charname,
				conn.RemoteAddr(),
				r)
		}
	}()

	for {
		if user == nil {
			menu := mainMenu()
			choice, _ := menu.Exec(conn, utils.ColorModeNone)

			switch choice {
			case "l":
				user = login(conn)
			case "n":
				user = newUser(conn)
			case "":
				fallthrough
			case "q":
				quit(conn)
				return
			}
		} else if character == nil {
			menu := userMenu(user)
			choice, charId := menu.Exec(conn, user.GetColorMode())

			switch choice {
			case "":
				fallthrough
			case "l":
				user.SetOnline(false)
				user = nil
			case "a":
				adminMenu := adminMenu()
				for {
					choice, _ := adminMenu.Exec(conn, user.GetColorMode())
					if choice == "" {
						break
					} else if choice == "u" {
						for {
							userAdminMenu := userAdminMenu()
							choice, userId := userAdminMenu.Exec(conn, user.GetColorMode())
							if choice == "" {
								break
							} else {
								_, err := strconv.Atoi(choice)

								if err == nil {
									for {
										userMenu := userSpecificMenu(model.M.GetUser(userId))
										choice, _ = userMenu.Exec(conn, user.GetColorMode())
										if choice == "" {
											break
										} else if choice == "d" {
											model.M.DeleteUser(userId)
											break
										}
									}
								}
							}
						}
					}
				}
			case "n":
				character = newCharacter(conn, user)
			case "d":
				for {
					deleteMenu := deleteMenu(user)
					deleteChoice, deleteCharId := deleteMenu.Exec(conn, user.GetColorMode())

					if deleteChoice == "" || deleteChoice == "c" {
						break
					}

					_, err := strconv.Atoi(deleteChoice)

					if err == nil {
						// TODO: Delete confirmation
						model.M.DeleteCharacter(deleteCharId)
					}
				}

			default:
				_, err := strconv.Atoi(choice)

				if err == nil {
					character = model.M.GetCharacter(charId)
				}
			}
		} else {
			game.Exec(conn, user, character)
			character = nil
		}
	}
}

func Exec() {
	fmt.Printf("Connecting to database... ")
	session, err := mgo.Dial("localhost")

	utils.HandleError(err)

	fmt.Println("done.")

	listener, err := net.Listen("tcp", ":8945")
	utils.HandleError(err)

	err = model.Init(session.Copy())

	// If there are no rooms at all create one
	rooms := model.M.GetRooms()
	if len(rooms) == 0 {
		zones := model.M.GetZones()

		var zone *database.Zone

		if len(zones) == 0 {
			zone = model.M.CreateZone("Default")
		} else {
			zone = zones[0]
		}

		model.M.CreateRoom(zone)
	}

	fmt.Println("Server listening on port 8945")

	for {
		conn, err := listener.Accept()
		utils.HandleError(err)
		fmt.Println("Client connected:", conn.RemoteAddr())
		go handleConnection(conn)
	}
}

// vim: nocindent
<|endoftext|>


<commit_before>package main

import (
    "encoding/binary"
    "errors"
    "flag"
    "fmt"
    ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
    "io"
    "log"
    "net"
    "os"
    "os/signal"
    "runtime"
    "strconv"
    "syscall"
)

var debug ss.DebugLog

var storage *Storage

const dnsGoroutineNum = 64

func getUser(conn net.Conn) (user User, err error) {
    const (
        idVersion = 0 // \x01
        version   = 1
        idUser    = 1 // \x02
    )
    buf := make([]byte, 2)
    if _, err = io.ReadFull(conn, buf); err != nil {
        return
    }
    switch buf[idVersion] {
    case version:
        break
    default:
        err = errors.New(fmt.Sprintf("version %s not supported", buf[idVersion]))
        return
    }
    userLen := buf[idUser]
    username := make([]byte, userLen)
    if _, err = io.ReadFull(conn, username); err != nil {
        return
    }
    user, err = storage.Get(string(username))
    return
}

func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
    const (
        idType  = 0 // address type index
        idIP0   = 1 // ip addres start index
        idDmLen = 1 // domain address length index
        idDm0   = 2 // domain address start index

        typeIPv4 = 1 // type is ipv4 address
        typeDm   = 3 // type is domain address
        typeIPv6 = 4 // type is ipv6 address

        lenIPv4   = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
        lenIPv6   = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
        lenDmBase = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
    )

    // buf size should at least have the same size with the largest possible
    // request size (when addrType is 3, domain name has at most 256 bytes)
    // 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
    buf := make([]byte, 260)
    var n int
    // read till we get possible domain length field
    ss.SetReadTimeout(conn)
    if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
        return
    }

    reqLen := -1
    switch buf[idType] {
    case typeIPv4:
        reqLen = lenIPv4
    case typeIPv6:
        reqLen = lenIPv6
    case typeDm:
        reqLen = int(buf[idDmLen]) + lenDmBase
    default:
        err = errors.New(fmt.Sprintf("addr type %d not supported", buf[idType]))
        return
    }

    if n < reqLen { // rare case
        ss.SetReadTimeout(conn)
        if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
            return
        }
    } else if n > reqLen {
        // it's possible to read more than just the request head
        extra = buf[reqLen:n]
    }

    // Return string for typeIP is not most efficient, but browsers (Chrome,
    // Safari, Firefox) all seems using typeDm exclusively. So this is not a
    // big problem.
    switch buf[idType] {
    case typeIPv4:
        host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
    case typeIPv6:
        host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
    case typeDm:
        host = string(buf[idDm0 : idDm0+buf[idDmLen]])
    }
    // parse port
    port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen])
    host = net.JoinHostPort(host, strconv.Itoa(int(port)))
    return
}

const logCntDelta = 100

var connCnt int
var nextLogConnCnt int = logCntDelta

func handleConnection(conn *ss.Conn) {
    var host string

    connCnt++ // this maybe not accurate, but should be enough
    if connCnt-nextLogConnCnt >= 0 {
        // XXX There's no xadd in the atomic package, so it's difficult to log
        // the message only once with low cost. Also note nextLogConnCnt maybe
        // added twice for current peak connection number level.
        log.Printf("Number of client connections reaches %d\n", nextLogConnCnt)
        nextLogConnCnt += logCntDelta
    }

    // function arguments are always evaluated, so surround debug statement
    // with if statement
    if debug {
        debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr())
    }
    closed := false
    defer func() {
        if debug {
            debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host)
        }
        connCnt--
        if !closed {
            conn.Close()
        }
    }()

    host, extra, err := getRequest(conn)
    if err != nil {
        log.Println("error getting request", conn.RemoteAddr(), conn.LocalAddr(), err)
        return
    }
    debug.Println("connecting", host)
    remote, err := net.Dial("tcp", host)
    if err != nil {
        if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
            // log too many open file error
            // EMFILE is process reaches open file limits, ENFILE is system limit
            log.Println("dial error:", err)
        } else {
            log.Println("error connecting to:", host, err)
        }
        return
    }
    defer func() {
        if !closed {
            remote.Close()
        }
    }()
    // write extra bytes read from
    if extra != nil {
        // debug.Println("getRequest read extra data, writing to remote, len", len(extra))
        if _, err = remote.Write(extra); err != nil {
            debug.Println("write request extra error:", err)
            return
        }
    }
    if debug {
        debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
    }
    go ss.PipeThenClose(conn, remote, ss.SET_TIMEOUT)
    ss.PipeThenClose(remote, conn, ss.NO_TIMEOUT)
    closed = true
    return
}

func waitSignal() {
    var sigChan = make(chan os.Signal, 1)
    signal.Notify(sigChan, syscall.SIGHUP)
    for sig := range sigChan {
        if sig == syscall.SIGHUP {
        } else {
            // is this going to happen?
            log.Printf("caught signal %v, exit", sig)
            os.Exit(0)
        }
    }
}

func run(port string) {
    ln, err := net.Listen("tcp", ":"+port)
    if err != nil {
        log.Printf("error listening port %v: %v\n", port, err)
        return
    }
    var cipher *ss.Cipher
    log.Printf("server listening port %v ...\n", port)
    for {
        conn, err := ln.Accept()
        if err != nil {
            debug.Printf("accept error: %v\n", err)
            return
        }

        user, err := getUser(conn)
        if err != nil {
            log.Printf("Error get passeoed: %s %v\n", port, err)
            conn.Close()
            continue
        }
        log.Println("creating cipher for user:", user.Name)
        cipher, err = ss.NewCipher(user.Method, user.Password)
        if err != nil {
            log.Printf("Error generating cipher for user: %s %v\n", user.Name, err)
            conn.Close()
            continue
        }
        go handleConnection(ss.NewConn(conn, cipher.Copy()))
    }
}

func main() {
    log.SetOutput(os.Stdout)

    var printVer bool
    var core int
    var serverPort string
    var redisServer string

    flag.BoolVar(&printVer, "version", false, "print version")
    flag.StringVar(&serverPort, "p", "8388", "server port")
    flag.StringVar(&redisServer, "redis", ":6379", "redis server")
    flag.IntVar(&core, "core", 0, "maximum number of CPU cores to use, default is determinied by Go runtime")
    flag.BoolVar((*bool)(&debug), "d", false, "print debug message")

    flag.Parse()

    if printVer {
        ss.PrintVersion()
        os.Exit(0)
    }

    ss.SetDebug(debug)

    storage = NewStorage(redisServer)
    if core > 0 {
        runtime.GOMAXPROCS(core)
    }

    go run(serverPort)

    waitSignal()
}
<commit_msg>Add use connect log<commit_after>package main

import (
    "bytes"
    "encoding/binary"
    "errors"
    "flag"
    "fmt"
    "github.com/cyfdecyf/leakybuf"
    ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
    "io"
    "log"
    "net"
    "os"
    "os/signal"
    "runtime"
    "strconv"
    "syscall"
)

var debug ss.DebugLog

var storage *Storage

const dnsGoroutineNum = 64

func getUser(conn net.Conn) (user User, err error) {
    const (
        idVersion = 0 // \x01
        version   = 1
        idUser    = 1 // \x02
    )
    buf := make([]byte, 2)
    if _, err = io.ReadFull(conn, buf); err != nil {
        return
    }
    switch buf[idVersion] {
    case version:
        break
    default:
        err = errors.New(fmt.Sprintf("version %s not supported", buf[idVersion]))
        return
    }
    userLen := buf[idUser]
    username := make([]byte, userLen)
    if _, err = io.ReadFull(conn, username); err != nil {
        return
    }
    user, err = storage.Get(string(username))
    return
}

func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
    const (
        idType  = 0 // address type index
        idIP0   = 1 // ip addres start index
        idDmLen = 1 // domain address length index
        idDm0   = 2 // domain address start index

        typeIPv4 = 1 // type is ipv4 address
        typeDm   = 3 // type is domain address
        typeIPv6 = 4 // type is ipv6 address

        lenIPv4   = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
        lenIPv6   = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
        lenDmBase = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
    )

    // buf size should at least have the same size with the largest possible
    // request size (when addrType is 3, domain name has at most 256 bytes)
    // 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
    buf := make([]byte, 260)
    var n int
    // read till we get possible domain length field
    ss.SetReadTimeout(conn)
    if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
        return
    }

    reqLen := -1
    switch buf[idType] {
    case typeIPv4:
        reqLen = lenIPv4
    case typeIPv6:
        reqLen = lenIPv6
    case typeDm:
        reqLen = int(buf[idDmLen]) + lenDmBase
    default:
        err = errors.New(fmt.Sprintf("addr type %d not supported", buf[idType]))
        return
    }

    if n < reqLen { // rare case
        ss.SetReadTimeout(conn)
        if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
            return
        }
    } else if n > reqLen {
        // it's possible to read more than just the request head
        extra = buf[reqLen:n]
    }

    // Return string for typeIP is not most efficient, but browsers (Chrome,
    // Safari, Firefox) all seems using typeDm exclusively. So this is not a
    // big problem.
    switch buf[idType] {
    case typeIPv4:
        host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
    case typeIPv6:
        host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
    case typeDm:
        host = string(buf[idDm0 : idDm0+buf[idDmLen]])
    }
    // parse port
    port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen])
    host = net.JoinHostPort(host, strconv.Itoa(int(port)))
    return
}

const logCntDelta = 100

var connCnt int
var nextLogConnCnt int = logCntDelta

func handleConnection(user User, conn *ss.Conn) {
    var host string
    var size = 0

    connCnt++ // this maybe not accurate, but should be enough
    if connCnt-nextLogConnCnt >= 0 {
        // XXX There's no xadd in the atomic package, so it's difficult to log
        // the message only once with low cost. Also note nextLogConnCnt maybe
        // added twice for current peak connection number level.
        log.Printf("Number of client connections reaches %d\n", nextLogConnCnt)
        nextLogConnCnt += logCntDelta
    }

    // function arguments are always evaluated, so surround debug statement
    // with if statement
    if debug {
        debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr())
    }
    closed := false
    defer func() {
        if debug {
            debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host)
        }
        connCnt--
        if !closed {
            conn.Close()
        }
    }()

    host, extra, err := getRequest(conn)
    lines := bytes.SplitN(extra, []byte("\r\n"), 2)
    var uri []byte

    if len(lines) == 2 {
        uri = lines[0]
    }

    if err != nil {
        log.Println("error getting request", conn.RemoteAddr(), conn.LocalAddr(), err)
        return
    }
    debug.Println("connecting", host)
    remote, err := net.Dial("tcp", host)
    if err != nil {
        if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
            // log too many open file error
            // EMFILE is process reaches open file limits, ENFILE is system limit
            log.Println("dial error:", err)
        } else {
            log.Println("error connecting to:", host, err)
        }
        return
    }
    defer func() {
        if !closed {
            remote.Close()
        }
    }()
    // write extra bytes read from
    if extra != nil {
        // debug.Println("getRequest read extra data, writing to remote, len", len(extra))
        if _, err = remote.Write(extra); err != nil {
            debug.Println("write request extra error:", err)
            return
        }
    }
    if debug {
        debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
    }
    go ss.PipeThenClose(conn, remote, ss.SET_TIMEOUT)
    res_size, status_lines := PipeThenClose(remote, conn, ss.NO_TIMEOUT)
    if uri == nil {
        fmt.Printf("CONNET %s %s %d\n", host, user.Name, size)
    } else {
        size += res_size
        lines = bytes.SplitN(uri, []byte(" "), 2)
        status := bytes.SplitN(status_lines, []byte(" "), 3)
        status_code := status[1]
        fmt.Printf("%s http://%s%s %s %s %d\n", lines[0], host, lines[1], status_code, user.Name, size)
    }
    closed = true
    return
}

const bufSize = 4096
const nBuf = 2048

func PipeThenClose(src, dst net.Conn, timeoutOpt int) (total int, line []byte) {
    var pipeBuf = leakybuf.NewLeakyBuf(nBuf, bufSize)
    defer dst.Close()
    buf := pipeBuf.Get()
    // defer pipeBuf.Put(buf)
    var first = true
    var size int
    for {
        if timeoutOpt == ss.SET_TIMEOUT {
            ss.SetReadTimeout(src)
        }
        n, err := src.Read(buf)
        if first {
            lines := bytes.SplitN(buf, []byte("\r\n"), 2)
            if len(lines) == 2 {
                line = lines[0]
            }
            first = false
        }
        // read may return EOF with n > 0
        // should always process n > 0 bytes before handling error
        if n > 0 {
            if size, err = dst.Write(buf[0:n]); err != nil {
                ss.Debug.Println("write:", err)
                break
            }
            total += size
        }
        if err != nil || n == 0 {
            // Always "use of closed network connection", but no easy way to
            // identify this specific error. So just leave the error along for now.
            // More info here: https://code.google.com/p/go/issues/detail?id=4373
            break
        }
    }
    return
}

func waitSignal() {
    var sigChan = make(chan os.Signal, 1)
    signal.Notify(sigChan, syscall.SIGHUP)
    for sig := range sigChan {
        if sig == syscall.SIGHUP {
        } else {
            // is this going to happen?
            log.Printf("caught signal %v, exit", sig)
            os.Exit(0)
        }
    }
}

func run(port string) {
    ln, err := net.Listen("tcp", ":"+port)
    if err != nil {
        log.Printf("error listening port %v: %v\n", port, err)
        return
    }
    var cipher *ss.Cipher
    log.Printf("server listening port %v ...\n", port)
    for {
        conn, err := ln.Accept()
        if err != nil {
            debug.Printf("accept error: %v\n", err)
            return
        }

        user, err := getUser(conn)
        if err != nil {
            log.Printf("Error get passeoed: %s %v\n", port, err)
            conn.Close()
            continue
        }
        log.Println("creating cipher for user:", user.Name)
        cipher, err = ss.NewCipher(user.Method, user.Password)
        if err != nil {
            log.Printf("Error generating cipher for user: %s %v\n", user.Name, err)
            conn.Close()
            continue
        }
        go handleConnection(user, ss.NewConn(conn, cipher.Copy()))
    }
}

func main() {
    log.SetOutput(os.Stdout)

    var printVer bool
    var core int
    var serverPort string
    var redisServer string

    flag.BoolVar(&printVer, "version", false, "print version")
    flag.StringVar(&serverPort, "p", "8388", "server port")
    flag.StringVar(&redisServer, "redis", ":6379", "redis server")
    flag.IntVar(&core, "core", 0, "maximum number of CPU cores to use, default is determinied by Go runtime")
    flag.BoolVar((*bool)(&debug), "d", false, "print debug message")

    flag.Parse()

    if printVer {
        ss.PrintVersion()
        os.Exit(0)
    }

    ss.SetDebug(debug)

    storage = NewStorage(redisServer)
    if core > 0 {
        runtime.GOMAXPROCS(core)
    }

    go run(serverPort)

    waitSignal()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	r "github.com/garyburd/redigo/redis"
	"github.com/go-martini/martini"
	"github.com/martini-contrib/render"

	"github.com/vieux/gocover.io/server/redis"
)

var (
	docker_socket = flag.String("s", "", "Dockerd socket (e.g., /var/run/docker.sock)")
	docker_addr   = flag.String("d", "", "Dockerd addr (e.g., 127.0.0.1:2375)")
	serveAddr     = flag.String("p", ":8080", "Address and port to serve")
	redisAddr     = flag.String("r", "127.0.0.1:6379", "redis address")
	redisPass     = flag.String("rp", "", "redis password")
)

func docker(repo, version string, pool *r.Pool) (int, string) {
	var (
		worker = "vieux/gocover"
		conn   = pool.Get()
	)

	defer conn.Close()

	if version != "" {
		worker = worker + ":" + version
	}

	if version == "" {
		if cached, fresh, err := redis.GetRepo(conn, repo); err != nil {
			return 500, err.Error()
		} else if fresh {
			return 200, string(cached)
		}
	}

	host := ""

	if *docker_socket != "" {
		host = "unix://" + *docker_socket
	} else if *docker_addr != "" {
		host = "tcp://" + *docker_addr
	} else {
		return 500, "cannot connect to docker daemon"
	}

	out, err := exec.Command("docker", "-H", host, "run", "--rm", "-a", "stdout", "-a", "stderr", worker, repo).CombinedOutput()
	if err != nil {
		if strings.Contains(string(out), "Unable to find image") {
			return 500, "go version '" + version + "' not found"
		}
		return 500, string(out)
	}
	re, err := regexp.Compile("\\")
	if err != nil {
		return 500, err.Error()
	}
	content := re.ReplaceAllString(string(out), "")
	content = strings.Replace(content, "background: black;", "background: #222222;", 2)

	content = strings.Replace(content, ".cov1 { color: rgb(128, 128, 128) }", ".cov1 { color: #52987D }", 2)
	content = strings.Replace(content, ".cov2 { color: rgb(128, 128, 128) }", ".cov2 { color: #4BA180 }", 2)
	content = strings.Replace(content, ".cov3 { color: rgb(128, 128, 128) }", ".cov3 { color: #44AA83 }", 2)
	content = strings.Replace(content, ".cov4 { color: rgb(128, 128, 128) }", ".cov4 { color: #3DB487 }", 2)
	content = strings.Replace(content, ".cov5 { color: rgb(128, 128, 128) }", ".cov5 { color: #36BD8A }", 2)
	content = strings.Replace(content, ".cov6 { color: rgb(128, 128, 128) }", ".cov6 { color: #2FC68D }", 2)
	content = strings.Replace(content, ".cov7 { color: rgb(128, 128, 128) }", ".cov7 { color: #28D091 }", 2)
	content = strings.Replace(content, ".cov8 { color: rgb(128, 128, 128) }", ".cov8 { color: #21D994 }", 2)
	content = strings.Replace(content, ".cov9 { color: rgb(128, 128, 128) }", ".cov9 { color: #1AE297 }", 2)

	content = strings.Replace(content, "\">"+repo, "\">", -1)

	re = regexp.MustCompile("-- cov:([0-9.]*) --")
	matches := re.FindStringSubmatch(content)
	if len(matches) == 2 {
		cov, err := strconv.ParseFloat(matches[1], 64)
		if err == nil {
			content = strings.Replace(content, "<select id=", fmt.Sprintf("%s%% | <select id=", int((cov-0.0001)/10), matches[1]), 1)
		}
		if version != "" {
			content = strings.Replace(content, "<select id=", fmt.Sprintf("%s | <select id=", version), 1)
		} else {
			redis.SetCache(conn, repo, content, matches[1])
		}
	} else if version != "" {
		content = strings.Replace(content, "<select id=", fmt.Sprintf("%s | <select id=", version), 1)
	} else {
		redis.SetCache(conn, repo, content, "-1")
	}
	if version == "" {
		redis.SetStats(conn, repo)
	}
	return 200, content
}

func main() {
	flag.Parse()

	var (
		m         = martini.Classic()
		pool, err = redis.NewPool("tcp", *redisAddr, *redisPass)
	)

	if err != nil {
		log.Fatalf("%v", err)
	}

	m.Use(martini.Static("static"))
	m.Use(render.Renderer(render.Options{Layout: "layout"}))
	m.Get("/about", func(r render.Render) {
		r.HTML(200, "about", map[string]interface{}{"about_active": "active"})
	})
	m.Get("/", func(r render.Render) {
		conn := pool.Get()
		defer conn.Close()

		top, err := redis.Top(conn, "top", 4)
		if err != nil {
			log.Println(err.Error())
		}
		last, err := redis.Top(conn, "last", 4)
		if err != nil {
			log.Println(err.Error())
		}
		r.HTML(200, "cover", map[string]interface{}{"top": top, "last": last, "cover_active": "active"})
	})
	m.Post("/_webhook", func(req *http.Request) (int, string) {

		github := struct {
			Repository struct {
				Full_Name string
			}
		}{}

		if err := json.NewDecoder(req.Body).Decode(&github); err != nil {
			return 500, err.Error()
		}

		go docker(github.Repository.Full_Name, "", pool)

		return 202, github.Repository.Full_Name
	})
	m.Get("/_badge/**", func(params martini.Params, r render.Render) {
		var (
			repo = params["_1"]
			conn = pool.Get()
		)

		defer conn.Close()
		if coverage, err := redis.GetCoverage(conn, repo); err != nil {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-error-lightgrey.svg?style=flat"))
		} else if coverage < 25.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-red.svg?style=flat", coverage))
		} else if coverage < 50.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-orange.svg?style=flat", coverage))
		} else if coverage < 75.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-green.svg?style=flat", coverage))
		} else {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-brightgreen.svg?style=flat", coverage))
		}

	})
	m.Get("/_cache/**", func(req *http.Request, params martini.Params) (int, string) {
		var (
			repo = params["_1"]
			conn = pool.Get()
		)
		defer conn.Close()

		if req.FormValue("version") == "" {
			if cached, _, err := redis.GetRepo(conn, repo); err != nil {
				return 500, err.Error()
			} else if cached != "" {
				redis.SetStats(conn, repo)
				return 200, string(cached)
			}
		}
		return 404, "No cached version of " + repo
	})
	m.Get("/_/**", func(req *http.Request, params martini.Params) (int, string) {
		return docker(params["_1"], req.FormValue("version"), pool)
	})
	m.Get("/**", func(req *http.Request, params martini.Params, r render.Render) {
		var (
			repo    = params["_1"]
			conn    = pool.Get()
			version = ""
		)
		defer conn.Close()

		if req.ParseForm() == nil {
			version = req.FormValue("version")
		}

		if cached, fresh, err := redis.GetRepo(conn, repo+version); err != nil {
			r.HTML(500, "", map[string]interface{}{"cover_active": "active", "error": err})
		} else if fresh {
			redis.SetStats(conn, repo)
			r.HTML(200, "cached", map[string]interface{}{"repo": repo, "cover_active": "active", "cache": template.HTML(cached), "version": version})
		} else {
			contexts := map[string]interface{}{"repo": repo, "cover_active": "active", "version": version}
			if cached != "" {
				contexts["cache"] = "ok"
			}
			r.HTML(200, "loading", contexts)
		}
	})
	log.Fatal(http.ListenAndServe(*serveAddr, m))
}
<commit_msg>support TLS<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	r "github.com/garyburd/redigo/redis"
	"github.com/go-martini/martini"
	"github.com/martini-contrib/render"

	"github.com/vieux/gocover.io/server/redis"
)

var (
	docker_socket = flag.String("s", "", "Dockerd socket (e.g., /var/run/docker.sock)")
	docker_addr   = flag.String("d", "", "Dockerd addr (e.g., 127.0.0.1:2375)")
	serveAddr     = flag.String("p", ":8080", "Address and port to serve")
	serveSAddr    = flag.String("ps", ":80443", "Address and port to serve HTTPS")
	redisAddr     = flag.String("r", "127.0.0.1:6379", "redis address")
	redisPass     = flag.String("rp", "", "redis password")
	certPath      = flag.String("tls", "", "cert path")
)

func docker(repo, version string, pool *r.Pool) (int, string) {
	var (
		worker = "vieux/gocover"
		conn   = pool.Get()
	)

	defer conn.Close()

	if version != "" {
		worker = worker + ":" + version
	}

	if version == "" {
		if cached, fresh, err := redis.GetRepo(conn, repo); err != nil {
			return 500, err.Error()
		} else if fresh {
			return 200, string(cached)
		}
	}

	host := ""

	if *docker_socket != "" {
		host = "unix://" + *docker_socket
	} else if *docker_addr != "" {
		host = "tcp://" + *docker_addr
	} else {
		return 500, "cannot connect to docker daemon"
	}

	out, err := exec.Command("docker", "-H", host, "run", "--rm", "-a", "stdout", "-a", "stderr", worker, repo).CombinedOutput()
	if err != nil {
		if strings.Contains(string(out), "Unable to find image") {
			return 500, "go version '" + version + "' not found"
		}
		return 500, string(out)
	}
	re, err := regexp.Compile("\\")
	if err != nil {
		return 500, err.Error()
	}
	content := re.ReplaceAllString(string(out), "")
	content = strings.Replace(content, "background: black;", "background: #222222;", 2)

	content = strings.Replace(content, ".cov1 { color: rgb(128, 128, 128) }", ".cov1 { color: #52987D }", 2)
	content = strings.Replace(content, ".cov2 { color: rgb(128, 128, 128) }", ".cov2 { color: #4BA180 }", 2)
	content = strings.Replace(content, ".cov3 { color: rgb(128, 128, 128) }", ".cov3 { color: #44AA83 }", 2)
	content = strings.Replace(content, ".cov4 { color: rgb(128, 128, 128) }", ".cov4 { color: #3DB487 }", 2)
	content = strings.Replace(content, ".cov5 { color: rgb(128, 128, 128) }", ".cov5 { color: #36BD8A }", 2)
	content = strings.Replace(content, ".cov6 { color: rgb(128, 128, 128) }", ".cov6 { color: #2FC68D }", 2)
	content = strings.Replace(content, ".cov7 { color: rgb(128, 128, 128) }", ".cov7 { color: #28D091 }", 2)
	content = strings.Replace(content, ".cov8 { color: rgb(128, 128, 128) }", ".cov8 { color: #21D994 }", 2)
	content = strings.Replace(content, ".cov9 { color: rgb(128, 128, 128) }", ".cov9 { color: #1AE297 }", 2)

	content = strings.Replace(content, "\">"+repo, "\">", -1)

	re = regexp.MustCompile("-- cov:([0-9.]*) --")
	matches := re.FindStringSubmatch(content)
	if len(matches) == 2 {
		cov, err := strconv.ParseFloat(matches[1], 64)
		if err == nil {
			content = strings.Replace(content, "<select id=", fmt.Sprintf("%s%% | <select id=", int((cov-0.0001)/10), matches[1]), 1)
		}
		if version != "" {
			content = strings.Replace(content, "<select id=", fmt.Sprintf("%s | <select id=", version), 1)
		} else {
			redis.SetCache(conn, repo, content, matches[1])
		}
	} else if version != "" {
		content = strings.Replace(content, "<select id=", fmt.Sprintf("%s | <select id=", version), 1)
	} else {
		redis.SetCache(conn, repo, content, "-1")
	}
	if version == "" {
		redis.SetStats(conn, repo)
	}
	return 200, content
}

func main() {
	flag.Parse()

	var (
		m         = martini.Classic()
		pool, err = redis.NewPool("tcp", *redisAddr, *redisPass)
	)

	if err != nil {
		log.Fatalf("%v", err)
	}

	m.Use(martini.Static("static"))
	m.Use(render.Renderer(render.Options{Layout: "layout"}))
	m.Get("/about", func(r render.Render) {
		r.HTML(200, "about", map[string]interface{}{"about_active": "active"})
	})
	m.Get("/", func(r render.Render) {
		conn := pool.Get()
		defer conn.Close()

		top, err := redis.Top(conn, "top", 4)
		if err != nil {
			log.Println(err.Error())
		}
		last, err := redis.Top(conn, "last", 4)
		if err != nil {
			log.Println(err.Error())
		}
		r.HTML(200, "cover", map[string]interface{}{"top": top, "last": last, "cover_active": "active"})
	})
	m.Post("/_webhook", func(req *http.Request) (int, string) {

		github := struct {
			Repository struct {
				Full_Name string
			}
		}{}

		if err := json.NewDecoder(req.Body).Decode(&github); err != nil {
			return 500, err.Error()
		}

		go docker(github.Repository.Full_Name, "", pool)

		return 202, github.Repository.Full_Name
	})
	m.Get("/_badge/**", func(params martini.Params, r render.Render) {
		var (
			repo = params["_1"]
			conn = pool.Get()
		)

		defer conn.Close()
		if coverage, err := redis.GetCoverage(conn, repo); err != nil {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-error-lightgrey.svg?style=flat"))
		} else if coverage < 25.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-red.svg?style=flat", coverage))
		} else if coverage < 50.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-orange.svg?style=flat", coverage))
		} else if coverage < 75.0 {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-green.svg?style=flat", coverage))
		} else {
			r.Redirect(fmt.Sprintf("https://img.shields.io/badge/coverage-%.1f%%-brightgreen.svg?style=flat", coverage))
		}

	})
	m.Get("/_cache/**", func(req *http.Request, params martini.Params) (int, string) {
		var (
			repo = params["_1"]
			conn = pool.Get()
		)
		defer conn.Close()

		if req.FormValue("version") == "" {
			if cached, _, err := redis.GetRepo(conn, repo); err != nil {
				return 500, err.Error()
			} else if cached != "" {
				redis.SetStats(conn, repo)
				return 200, string(cached)
			}
		}
		return 404, "No cached version of " + repo
	})
	m.Get("/_/**", func(req *http.Request, params martini.Params) (int, string) {
		return docker(params["_1"], req.FormValue("version"), pool)
	})
	m.Get("/**", func(req *http.Request, params martini.Params, r render.Render) {
		var (
			repo    = params["_1"]
			conn    = pool.Get()
			version = ""
		)
		defer conn.Close()

		if req.ParseForm() == nil {
			version = req.FormValue("version")
		}

		if cached, fresh, err := redis.GetRepo(conn, repo+version); err != nil {
			r.HTML(500, "", map[string]interface{}{"cover_active": "active", "error": err})
		} else if fresh {
			redis.SetStats(conn, repo)
			r.HTML(200, "cached", map[string]interface{}{"repo": repo, "cover_active": "active", "cache": template.HTML(cached), "version": version})
		} else {
			contexts := map[string]interface{}{"repo": repo, "cover_active": "active", "version": version}
			if cached != "" {
				contexts["cache"] = "ok"
			}
			r.HTML(200, "loading", contexts)
		}
	})
	if *certPath != "" {
		go func() {
			log.Println(http.ListenAndServe(*serveAddr, http.RedirectHandler("https://gocover.io", http.StatusMovedPermanently)))
		}()
		log.Fatal(http.ListenAndServeTLS(*serveSAddr, filepath.Join(*certPath, "cert.pem"), filepath.Join(*certPath, "privkey.pem"), m))
	} else {
		log.Fatal(http.ListenAndServe(*serveAddr, m))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema1"
	humanize "github.com/dustin/go-humanize"
	"github.com/jessfraz/reg/clair"
	"github.com/jessfraz/reg/registry"
	"github.com/jessfraz/reg/utils"
	"github.com/urfave/cli"
)

const (
	// VERSION is the binary version.
	VERSION = "v0.1.0"

	dockerConfigPath = ".docker/config.json"
)

var (
	updating = false
	wg       sync.WaitGroup
)

// preload initializes any global options and configuration
// before the main or sub commands are run.
func preload(c *cli.Context) (err error) {
	if c.GlobalBool("debug") {
		logrus.SetLevel(logrus.DebugLevel)
	}

	return nil
}

func main() {
	app := cli.NewApp()
	app.Name = "reg-server"
	app.Version = VERSION
	app.Author = "@jessfraz"
	app.Email = "no-reply@butts.com"
	app.Usage = "Docker registry v2 static UI server."
	app.Before = preload
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "run in debug mode",
		},
		cli.StringFlag{
			Name:  "username, u",
			Usage: "username for the registry",
		},
		cli.StringFlag{
			Name:  "password, p",
			Usage: "password for the registry",
		},
		cli.StringFlag{
			Name:  "registry, r",
			Usage: "URL to the provate registry (ex. r.j3ss.co)",
		},
		cli.StringFlag{
			Name:  "port",
			Value: "8080",
			Usage: "port for server to run on",
		},
		cli.StringFlag{
			Name:  "cert",
			Usage: "path to ssl cert",
		},
		cli.StringFlag{
			Name:  "key",
			Usage: "path to ssl key",
		},
		cli.StringFlag{
			Name:  "interval",
			Value: "5m",
			Usage: "interval to generate new index.html's at",
		},
		cli.StringFlag{
			Name:  "clair",
			Usage: "url to clair instance",
		},
	}
	app.Action = func(c *cli.Context) error {
		auth, err := utils.GetAuthConfig(c)
		if err != nil {
			return err
		}

		// create the registry client
		r, err := registry.New(auth, c.GlobalBool("debug"))
		if err != nil {
			return err
		}

		// get the path to the static directory
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		staticDir := filepath.Join(wd, "static")

		// create the initial index
		if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
			return err
		}

		// parse the duration
		dur, err := time.ParseDuration(c.String("interval"))
		if err != nil {
			logrus.Fatalf("parsing %s as duration failed: %v", c.String("interval"), err)
		}
		ticker := time.NewTicker(dur)

		go func() {
			// create more indexes every X minutes based off interval
			for range ticker.C {
				if !updating {
					if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
						logrus.Warnf("creating static index failed: %v", err)
						wg.Wait()
						updating = false
					}
				}
			}
		}()

		// create mux server
		mux := http.NewServeMux()

		// static files handler
		staticHandler := http.FileServer(http.Dir(staticDir))
		mux.Handle("/", staticHandler)

		// set up the server
		port := c.String("port")
		server := &http.Server{
			Addr:    ":" + port,
			Handler: mux,
		}
		logrus.Infof("Starting server on port %q", port)
		if c.String("cert") != "" && c.String("key") != "" {
			logrus.Fatal(server.ListenAndServeTLS(c.String("cert"), c.String("key")))
		} else {
			logrus.Fatal(server.ListenAndServe())
		}

		return nil
	}

	app.Run(os.Args)
}

type data struct {
	RegistryURL string
	LastUpdated string
	Repos       []repository
}

type repository struct {
	Name        string
	Tag         string
	RepoURI     string
	CreatedDate string
	VulnURI     string
}

type v1Compatibility struct {
	ID      string    `json:"id"`
	Created time.Time `json:"created"`
}

func createStaticIndex(r *registry.Registry, staticDir, clairURI string) error {
	updating = true
	logrus.Info("fetching catalog")
	repoList, err := r.Catalog("")
	if err != nil {
		return fmt.Errorf("getting catalog failed: %v", err)
	}

	logrus.Info("fetching tags")
	var repos []repository
	for i, repo := range repoList {
		// get the tags
		tags, err := r.Tags(repo)
		if err != nil {
			return fmt.Errorf("getting tags for %s failed: %v", repo, err)
		}
		for j, tag := range tags {
			// get the manifest

			manifest, err := r.Manifest(repo, tag)
			if err != nil {
				return fmt.Errorf("getting tags for %s:%s failed: %v", repo, tag, err)
			}

			var createdDate string
			if m1, ok := manifest.(schema1.SignedManifest); ok {
				history := m1.History
				for _, h := range history {
					var comp v1Compatibility
					if err := json.Unmarshal([]byte(h.V1Compatibility), &comp); err != nil {
						return fmt.Errorf("unmarshal v1compatibility failed: %v", err)
					}
					createdDate = humanize.Time(comp.Created)
					break
				}
			}

			repoURI := fmt.Sprintf("%s/%s", r.Domain, repo)
			if tag != "latest" {
				repoURI += ":" + tag
			}

			newrepo := repository{
				Name:        repo,
				Tag:         tag,
				RepoURI:     repoURI,
				CreatedDate: createdDate,
			}

			if clairURI != "" {
				wg.Add(1)

				go func(repo, tag string, i, j int) {
					defer wg.Done()

					throttle := time.Tick(time.Duration(time.Duration((i+1)*(j+1)*2) * time.Second))
					<-throttle

					logrus.Infof("creating vulns.txt for %s:%s", repo, tag)

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag); err != nil {
						// return fmt.Errorf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
						logrus.Warnf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
					}
				}(repo, tag, i, j)

				newrepo.VulnURI = filepath.Join(repo, tag, "vulns.txt")
			}
			repos = append(repos, newrepo)
		}
	}

	// create temporoary file to save template to
	logrus.Info("creating temporary file for template")
	f, err := ioutil.TempFile("", "reg-server")
	if err != nil {
		return fmt.Errorf("creating temp file failed: %v", err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// parse & execute the template
	logrus.Info("parsing and executing the template")
	templateDir := filepath.Join(staticDir, "../templates")
	lp := filepath.Join(templateDir, "layout.html")

	d := data{
		RegistryURL: r.Domain,
		Repos:       repos,
		LastUpdated: time.Now().Local().Format(time.RFC1123),
	}
	tmpl := template.Must(template.New("").ParseFiles(lp))
	if err := tmpl.ExecuteTemplate(f, "layout", d); err != nil {
		return fmt.Errorf("execute template failed: %v", err)
	}
	f.Close()

	index := filepath.Join(staticDir, "index.html")
	logrus.Infof("renaming the temporary file %s to %s", f.Name(), index)
	if err := os.Rename(f.Name(), index); err != nil {
		return fmt.Errorf("renaming result from %s to %s failed: %v", f.Name(), index, err)
	}
	updating = false
	return nil
}

func createVulnStaticPage(r *registry.Registry, staticDir, clairURI, repo, tag string) error {
	// get the manifest
	m, err := r.ManifestV1(repo, tag)
	if err != nil {
		return err
	}

	// filter out the empty layers
	var filteredLayers []schema1.FSLayer
	for _, layer := range m.FSLayers {
		if layer.BlobSum != clair.EmptyLayerBlobSum {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	m.FSLayers = filteredLayers
	if len(m.FSLayers) == 0 {
		fmt.Printf("No need to analyse image %s:%s as there is no non-emtpy layer", repo, tag)
		return nil
	}

	// initialize clair
	cr, err := clair.New(clairURI, false)
	if err != nil {
		return err
	}

	for i := len(m.FSLayers) - 1; i >= 0; i-- {
		// form the clair layer
		l, err := utils.NewClairLayer(r, repo, m.FSLayers, i)
		if err != nil {
			return err
		}

		// post the layer
		if _, err := cr.PostLayer(l); err != nil {
			return err
		}
	}

	vl, err := cr.GetLayer(m.FSLayers[0].BlobSum.String(), false, true)
	if err != nil {
		return err
	}

	// get the vulns
	var vulns []clair.Vulnerability
	for _, f := range vl.Features {
		for _, v := range f.Vulnerabilities {
			vulns = append(vulns, v)
		}
	}

	path := filepath.Join(staticDir, repo, tag, "vulns.txt")
	if err := os.MkdirAll(filepath.Dir(path), 0644); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	fmt.Fprintf(file, "Found %d vulnerabilities \n", len(vulns))
	fmt.Fprintln(file, "")

	vulnsBy := func(sev string, store map[string][]clair.Vulnerability) []clair.Vulnerability {
		items, found := store[sev]
		if !found {
			items = make([]clair.Vulnerability, 0)
			store[sev] = items
		}
		return items
	}

	// group by severity
	store := make(map[string][]clair.Vulnerability)
	for _, v := range vulns {
		sevRow := vulnsBy(v.Severity, store)
		store[v.Severity] = append(sevRow, v)
	}

	// iterate over the priorities list
	iteratePriorities := func(f func(sev string)) {
		for _, sev := range clair.Priorities {
			if len(store[sev]) != 0 {
				f(sev)
			}
		}
	}

	iteratePriorities(func(sev string) {
		fmt.Fprintf(file, "%s: %d\n", sev, len(store[sev]))
	})
	fmt.Fprintln(file, "")

	// return an error if there are more than 10 bad vulns
	lenBadVulns := len(store["High"]) + len(store["Critical"]) + len(store["Defcon1"])
	if lenBadVulns > 10 {
		fmt.Fprintln(file, "--------------- ALERT ---------------")
		fmt.Fprintf(file, "%d bad vunerabilities found", lenBadVulns)
	}
	fmt.Fprintln(file, "")

	iteratePriorities(func(sev string) {
		for _, v := range store[sev] {
			fmt.Fprintf(file, "%s: [%s] \n%s\n%s\n", v.Name, v.Severity, v.Description, v.Link)
			fmt.Fprintln(file, "-----------------------------------------")
		}
	})

	return nil
}
<commit_msg>longer<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema1"
	humanize "github.com/dustin/go-humanize"
	"github.com/jessfraz/reg/clair"
	"github.com/jessfraz/reg/registry"
	"github.com/jessfraz/reg/utils"
	"github.com/urfave/cli"
)

const (
	// VERSION is the binary version.
	VERSION = "v0.1.0"

	dockerConfigPath = ".docker/config.json"
)

var (
	updating = false
	wg       sync.WaitGroup
)

// preload initializes any global options and configuration
// before the main or sub commands are run.
func preload(c *cli.Context) (err error) {
	if c.GlobalBool("debug") {
		logrus.SetLevel(logrus.DebugLevel)
	}

	return nil
}

func main() {
	app := cli.NewApp()
	app.Name = "reg-server"
	app.Version = VERSION
	app.Author = "@jessfraz"
	app.Email = "no-reply@butts.com"
	app.Usage = "Docker registry v2 static UI server."
	app.Before = preload
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "run in debug mode",
		},
		cli.StringFlag{
			Name:  "username, u",
			Usage: "username for the registry",
		},
		cli.StringFlag{
			Name:  "password, p",
			Usage: "password for the registry",
		},
		cli.StringFlag{
			Name:  "registry, r",
			Usage: "URL to the provate registry (ex. r.j3ss.co)",
		},
		cli.StringFlag{
			Name:  "port",
			Value: "8080",
			Usage: "port for server to run on",
		},
		cli.StringFlag{
			Name:  "cert",
			Usage: "path to ssl cert",
		},
		cli.StringFlag{
			Name:  "key",
			Usage: "path to ssl key",
		},
		cli.StringFlag{
			Name:  "interval",
			Value: "5m",
			Usage: "interval to generate new index.html's at",
		},
		cli.StringFlag{
			Name:  "clair",
			Usage: "url to clair instance",
		},
	}
	app.Action = func(c *cli.Context) error {
		auth, err := utils.GetAuthConfig(c)
		if err != nil {
			return err
		}

		// create the registry client
		r, err := registry.New(auth, c.GlobalBool("debug"))
		if err != nil {
			return err
		}

		// get the path to the static directory
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		staticDir := filepath.Join(wd, "static")

		// create the initial index
		if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
			return err
		}

		// parse the duration
		dur, err := time.ParseDuration(c.String("interval"))
		if err != nil {
			logrus.Fatalf("parsing %s as duration failed: %v", c.String("interval"), err)
		}
		ticker := time.NewTicker(dur)

		go func() {
			// create more indexes every X minutes based off interval
			for range ticker.C {
				if !updating {
					if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
						logrus.Warnf("creating static index failed: %v", err)
						wg.Wait()
						updating = false
					}
				}
			}
		}()

		// create mux server
		mux := http.NewServeMux()

		// static files handler
		staticHandler := http.FileServer(http.Dir(staticDir))
		mux.Handle("/", staticHandler)

		// set up the server
		port := c.String("port")
		server := &http.Server{
			Addr:    ":" + port,
			Handler: mux,
		}
		logrus.Infof("Starting server on port %q", port)
		if c.String("cert") != "" && c.String("key") != "" {
			logrus.Fatal(server.ListenAndServeTLS(c.String("cert"), c.String("key")))
		} else {
			logrus.Fatal(server.ListenAndServe())
		}

		return nil
	}

	app.Run(os.Args)
}

type data struct {
	RegistryURL string
	LastUpdated string
	Repos       []repository
}

type repository struct {
	Name        string
	Tag         string
	RepoURI     string
	CreatedDate string
	VulnURI     string
}

type v1Compatibility struct {
	ID      string    `json:"id"`
	Created time.Time `json:"created"`
}

func createStaticIndex(r *registry.Registry, staticDir, clairURI string) error {
	updating = true
	logrus.Info("fetching catalog")
	repoList, err := r.Catalog("")
	if err != nil {
		return fmt.Errorf("getting catalog failed: %v", err)
	}

	logrus.Info("fetching tags")
	var repos []repository
	for i, repo := range repoList {
		// get the tags
		tags, err := r.Tags(repo)
		if err != nil {
			return fmt.Errorf("getting tags for %s failed: %v", repo, err)
		}
		for j, tag := range tags {
			// get the manifest

			manifest, err := r.Manifest(repo, tag)
			if err != nil {
				return fmt.Errorf("getting tags for %s:%s failed: %v", repo, tag, err)
			}

			var createdDate string
			if m1, ok := manifest.(schema1.SignedManifest); ok {
				history := m1.History
				for _, h := range history {
					var comp v1Compatibility
					if err := json.Unmarshal([]byte(h.V1Compatibility), &comp); err != nil {
						return fmt.Errorf("unmarshal v1compatibility failed: %v", err)
					}
					createdDate = humanize.Time(comp.Created)
					break
				}
			}

			repoURI := fmt.Sprintf("%s/%s", r.Domain, repo)
			if tag != "latest" {
				repoURI += ":" + tag
			}

			newrepo := repository{
				Name:        repo,
				Tag:         tag,
				RepoURI:     repoURI,
				CreatedDate: createdDate,
			}

			if clairURI != "" {
				wg.Add(1)

				go func(repo, tag string, i, j int) {
					defer wg.Done()

					throttle := time.Tick(time.Duration(time.Duration((i+1)*(j+1)*3) * time.Second))
					<-throttle

					logrus.Infof("creating vulns.txt for %s:%s", repo, tag)

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag); err != nil {
						// return fmt.Errorf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
						logrus.Warnf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
					}
				}(repo, tag, i, j)

				newrepo.VulnURI = filepath.Join(repo, tag, "vulns.txt")
			}
			repos = append(repos, newrepo)
		}
	}

	// create temporoary file to save template to
	logrus.Info("creating temporary file for template")
	f, err := ioutil.TempFile("", "reg-server")
	if err != nil {
		return fmt.Errorf("creating temp file failed: %v", err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// parse & execute the template
	logrus.Info("parsing and executing the template")
	templateDir := filepath.Join(staticDir, "../templates")
	lp := filepath.Join(templateDir, "layout.html")

	d := data{
		RegistryURL: r.Domain,
		Repos:       repos,
		LastUpdated: time.Now().Local().Format(time.RFC1123),
	}
	tmpl := template.Must(template.New("").ParseFiles(lp))
	if err := tmpl.ExecuteTemplate(f, "layout", d); err != nil {
		return fmt.Errorf("execute template failed: %v", err)
	}
	f.Close()

	index := filepath.Join(staticDir, "index.html")
	logrus.Infof("renaming the temporary file %s to %s", f.Name(), index)
	if err := os.Rename(f.Name(), index); err != nil {
		return fmt.Errorf("renaming result from %s to %s failed: %v", f.Name(), index, err)
	}
	updating = false
	return nil
}

func createVulnStaticPage(r *registry.Registry, staticDir, clairURI, repo, tag string) error {
	// get the manifest
	m, err := r.ManifestV1(repo, tag)
	if err != nil {
		return err
	}

	// filter out the empty layers
	var filteredLayers []schema1.FSLayer
	for _, layer := range m.FSLayers {
		if layer.BlobSum != clair.EmptyLayerBlobSum {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	m.FSLayers = filteredLayers
	if len(m.FSLayers) == 0 {
		fmt.Printf("No need to analyse image %s:%s as there is no non-emtpy layer", repo, tag)
		return nil
	}

	// initialize clair
	cr, err := clair.New(clairURI, false)
	if err != nil {
		return err
	}

	for i := len(m.FSLayers) - 1; i >= 0; i-- {
		// form the clair layer
		l, err := utils.NewClairLayer(r, repo, m.FSLayers, i)
		if err != nil {
			return err
		}

		// post the layer
		if _, err := cr.PostLayer(l); err != nil {
			return err
		}
	}

	vl, err := cr.GetLayer(m.FSLayers[0].BlobSum.String(), false, true)
	if err != nil {
		return err
	}

	// get the vulns
	var vulns []clair.Vulnerability
	for _, f := range vl.Features {
		for _, v := range f.Vulnerabilities {
			vulns = append(vulns, v)
		}
	}

	path := filepath.Join(staticDir, repo, tag, "vulns.txt")
	if err := os.MkdirAll(filepath.Dir(path), 0644); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	fmt.Fprintf(file, "Found %d vulnerabilities \n", len(vulns))
	fmt.Fprintln(file, "")

	vulnsBy := func(sev string, store map[string][]clair.Vulnerability) []clair.Vulnerability {
		items, found := store[sev]
		if !found {
			items = make([]clair.Vulnerability, 0)
			store[sev] = items
		}
		return items
	}

	// group by severity
	store := make(map[string][]clair.Vulnerability)
	for _, v := range vulns {
		sevRow := vulnsBy(v.Severity, store)
		store[v.Severity] = append(sevRow, v)
	}

	// iterate over the priorities list
	iteratePriorities := func(f func(sev string)) {
		for _, sev := range clair.Priorities {
			if len(store[sev]) != 0 {
				f(sev)
			}
		}
	}

	iteratePriorities(func(sev string) {
		fmt.Fprintf(file, "%s: %d\n", sev, len(store[sev]))
	})
	fmt.Fprintln(file, "")

	// return an error if there are more than 10 bad vulns
	lenBadVulns := len(store["High"]) + len(store["Critical"]) + len(store["Defcon1"])
	if lenBadVulns > 10 {
		fmt.Fprintln(file, "--------------- ALERT ---------------")
		fmt.Fprintf(file, "%d bad vunerabilities found", lenBadVulns)
	}
	fmt.Fprintln(file, "")

	iteratePriorities(func(sev string) {
		for _, v := range store[sev] {
			fmt.Fprintf(file, "%s: [%s] \n%s\n%s\n", v.Name, v.Severity, v.Description, v.Link)
			fmt.Fprintln(file, "-----------------------------------------")
		}
	})

	return nil
}
<|endoftext|>