<commit_before>package ipvs

import (
    "bytes"
    "encoding/hex"
    "net"
    "github.com/hkwi/nlgo"
    "syscall"
    "testing"
)

var testVersion = []struct { raw uint32; str string } {
    { 0x00010203, "1.2.3" },
}

func TestVersion (t *testing.T) {
    for _, test := range testVersion {
        ver := Version(test.raw)
        str := ver.String()

        if str != test.str {
            t.Errorf("fail %08x: %s != %s", test.raw, str, test.str)
        }
    }
}

func TestInfo (t *testing.T) {
    var info Info
    attrs := nlgo.AttrMap{Policy: ipvs_info_policy, AttrSlice: nlgo.AttrSlice{
        {Header: syscall.NlAttr{Type: IPVS_INFO_ATTR_VERSION}, Value: nlgo.U32(0x00010203)},
        {Header: syscall.NlAttr{Type: IPVS_INFO_ATTR_CONN_TAB_SIZE}, Value: nlgo.U32(4096)},
    }}

    if err := info.unpack(attrs); err != nil {
        t.Errorf("error Info.unpack(): %s", err)
    }

    if info.Version.String() != "1.2.3" {
        t.Errorf("fail Info.Version: %s != 1.2.3", info.Version.String())
    }

    if info.ConnTabSize != 4096 {
        t.Errorf("fail Info.ConnTabSize: %s != 4096", info.ConnTabSize)
    }
}

func TestService (t *testing.T) {
    var service Service

    pkt := []byte{
         0x06,0x00, 0x01,0x00, // IPVS_SVC_ATTR_AF
            0x02,0x00, 0x00,0x00, // 2
         0x06,0x00, 0x02,0x00, // IPVS_SVC_ATTR_PROTOCOL
            0x06,0x00, 0x00,0x00, // 6
         0x08,0x00, 0x03,0x00,   0x0a,0x6b,0x6b,0x00,   // IPVS_SVC_ATTR_ADDR       10.107.107.0
         0x06,0x00, 0x04,0x00,   0x05,0x39, 0x00,0x00,   // IPVS_SVC_ATTR_PORT       1337
         0x08,0x00, 0x06,0x00,   'w','l','c',0x00,      // IPVS_SVC_ATTR_SCHED_NAME wlc
         0x0c,0x00, 0x07,0x00,   0x00,0x00,0x00,0x00, 0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_FLAGS 0:0
         0x08,0x00, 0x08,0x00,   0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_TIMEOUT    0
         0x08,0x00, 0x09,0x00,   0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_NETMASK    0
    }
    if attrs, err := ipvs_service_policy.Parse(pkt); err != nil {
        t.Fatalf("error ipvs_service_policy.Parse: %s", err)
    } else if err := service.unpack(attrs.(nlgo.AttrMap)); err != nil {
        t.Fatalf("error Service.unpack: %s", err)
    }

    if service.Af != 2 {
        t.Errorf("fail Service.Af: %s", service.Af)
    }
    if service.Protocol != 6 {
        t.Errorf("fail Service.Protocol: %s", service.Protocol)
    }
    if service.Addr.String() != "10.107.107.0" {
        t.Errorf("fail Service.Addr: %s", service.Addr)
    }
    if service.Port != 1337 {
        t.Errorf("fail Service.Port: %s", service.Port)
    }
    if service.SchedName != "wlc" {
        t.Errorf("fail Service.SchedName: %s", service.SchedName)
    }
    if service.Flags.Flags != 0 || service.Flags.Mask != 0 {
        t.Errorf("fail Service.Flags: %+v", service.Flags)
    }
    if service.Timeout != 0 {
        t.Errorf("fail Service.Timeout: %s", service.Timeout)
    }
    if service.Netmask != 0 {
        t.Errorf("fail Service.Netmask: %s", service.Netmask)
    }

    outAttrs := service.attrs(true)
    outPkt := outAttrs.Bytes()

    if !bytes.Equal(outPkt, pkt) {
        t.Errorf("error Service.attrs: \n%s", hex.Dump(outPkt))
    }
}

func testDestEquals (t *testing.T, testDest Dest, dest Dest) {
    if dest.Addr.String() != testDest.Addr.String() {
        t.Errorf("fail testDest.unpack(): Addr %v", dest.Addr.String())
    }
    if dest.Port != testDest.Port {
        t.Errorf("fail testDest.unpack(): Port %v", dest.Port)
    }
    if dest.FwdMethod != testDest.FwdMethod {
        t.Errorf("fail testDest.unpack(): FwdMethod %v", dest.FwdMethod)
    }
    if dest.Weight != testDest.Weight {
        t.Errorf("fail testDest.unpack(): Weight %v", dest.Weight)
    }
    if dest.UThresh != testDest.UThresh {
        t.Errorf("fail testDest.unpack(): UThresh %v", dest.UThresh)
    }
    if dest.LThresh != testDest.LThresh {
        t.Errorf("fail testDest.unpack(): LThresh %v", dest.LThresh)
    }
}

func TestDest (t *testing.T) {
    testService := Service {
        Af:     syscall.AF_INET,
    }
    testDest := Dest{
        Addr:   net.ParseIP("10.107.107.0"),
        Port:   1337,

        FwdMethod:  IP_VS_CONN_F_TUNNEL,
        Weight:     10,
        UThresh:    1000,
        LThresh:    0,
    }
    testAttrs := nlgo.AttrSlice{
        nlattr(IPVS_DEST_ATTR_ADDR, nlgo.Binary([]byte{0x0a, 0x6b, 0x6b, 0x00})),
        nlattr(IPVS_DEST_ATTR_PORT, nlgo.U16(0x3905)),
        nlattr(IPVS_DEST_ATTR_FWD_METHOD, nlgo.U32(IP_VS_CONN_F_TUNNEL)),
        nlattr(IPVS_DEST_ATTR_WEIGHT, nlgo.U32(10)),
        nlattr(IPVS_DEST_ATTR_U_THRESH, nlgo.U32(1000)),
        nlattr(IPVS_DEST_ATTR_L_THRESH, nlgo.U32(0)),
    }

    // pack
    packAttrs := testDest.attrs(&testService, true)
    packBytes := packAttrs.Bytes()

    if !bytes.Equal(packBytes, testAttrs.Bytes()) {
        t.Errorf("fail Dest.attrs(): \n%s", hex.Dump(packBytes))
    }

    // unpack
    var unpackDest Dest

    if unpackAttrs, err := ipvs_dest_policy.Parse(packBytes); err != nil {
        t.Fatalf("error ipvs_dest_policy.Parse: %s", err)
    } else if err := unpackDest.unpack(testService, unpackAttrs.(nlgo.AttrMap)); err != nil {
        t.Fatalf("error Service.unpack: %s", err)
    }

    testDestEquals(t, testDest, unpackDest)
}
<commit_msg>ipvs: refactor TestService to pack-unpack<commit_after>package ipvs

import (
    "bytes"
    "encoding/hex"
    "net"
    "github.com/hkwi/nlgo"
    "syscall"
    "testing"
)

var testVersion = []struct { raw uint32; str string } {
    { 0x00010203, "1.2.3" },
}

func TestVersion (t *testing.T) {
    for _, test := range testVersion {
        ver := Version(test.raw)
        str := ver.String()

        if str != test.str {
            t.Errorf("fail %08x: %s != %s", test.raw, str, test.str)
        }
    }
}

func TestInfo (t *testing.T) {
    var info Info
    attrs := nlgo.AttrMap{Policy: ipvs_info_policy, AttrSlice: nlgo.AttrSlice{
        {Header: syscall.NlAttr{Type: IPVS_INFO_ATTR_VERSION}, Value: nlgo.U32(0x00010203)},
        {Header: syscall.NlAttr{Type: IPVS_INFO_ATTR_CONN_TAB_SIZE}, Value: nlgo.U32(4096)},
    }}

    if err := info.unpack(attrs); err != nil {
        t.Errorf("error Info.unpack(): %s", err)
    }

    if info.Version.String() != "1.2.3" {
        t.Errorf("fail Info.Version: %s != 1.2.3", info.Version.String())
    }

    if info.ConnTabSize != 4096 {
        t.Errorf("fail Info.ConnTabSize: %s != 4096", info.ConnTabSize)
    }
}

func testServiceEquals (t *testing.T, testService Service, service Service) {
    if service.Af != testService.Af {
        t.Errorf("fail Service.Af: %s", service.Af)
    }
    if service.Protocol != testService.Protocol {
        t.Errorf("fail Service.Protocol: %s", service.Protocol)
    }
    if service.Addr.String() != testService.Addr.String() {
        t.Errorf("fail Service.Addr: %s", service.Addr.String())
    }
    if service.Port != testService.Port {
        t.Errorf("fail Service.Port: %s", service.Port)
    }
    if service.SchedName != testService.SchedName {
        t.Errorf("fail Service.SchedName: %s", service.SchedName)
    }
    if service.Flags.Flags != testService.Flags.Flags || service.Flags.Mask != testService.Flags.Mask {
        t.Errorf("fail Service.Flags: %+v", service.Flags)
    }
    if service.Timeout != testService.Timeout {
        t.Errorf("fail Service.Timeout: %s", service.Timeout)
    }
    if service.Netmask != testService.Netmask {
        t.Errorf("fail Service.Netmask: %s", service.Netmask)
    }
}

func TestService (t *testing.T) {
    testService := Service {
        Af:         syscall.AF_INET,        // 2
        Protocol:   syscall.IPPROTO_TCP,    // 6
        Addr:       net.ParseIP("10.107.107.0"),
        Port:       1337,
        SchedName:  "wlc",
        Flags:      Flags{0, 0},
        Timeout:    0,
        Netmask:    0x00000000,
    }
    testBytes := []byte{
         0x06,0x00, 0x01,0x00, // IPVS_SVC_ATTR_AF
            0x02,0x00, 0x00,0x00, // 2
         0x06,0x00, 0x02,0x00, // IPVS_SVC_ATTR_PROTOCOL
            0x06,0x00, 0x00,0x00, // 6
         0x08,0x00, 0x03,0x00,   0x0a,0x6b,0x6b,0x00,   // IPVS_SVC_ATTR_ADDR       10.107.107.0
         0x06,0x00, 0x04,0x00,   0x05,0x39, 0x00,0x00,   // IPVS_SVC_ATTR_PORT       1337
         0x08,0x00, 0x06,0x00,   'w','l','c',0x00,      // IPVS_SVC_ATTR_SCHED_NAME wlc
         0x0c,0x00, 0x07,0x00,   0x00,0x00,0x00,0x00, 0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_FLAGS 0:0
         0x08,0x00, 0x08,0x00,   0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_TIMEOUT    0
         0x08,0x00, 0x09,0x00,   0x00,0x00,0x00,0x00,    // IPVS_SVC_ATTR_NETMASK    0
    }

    // pack
    packAttrs := testService.attrs(true)
    packBytes := packAttrs.Bytes()

    if !bytes.Equal(packBytes, testBytes) {
        t.Errorf("fail Dest.attrs(): \n%s", hex.Dump(packBytes))
    }

    // unpack
    var unpackService Service

    if unpackAttrs, err := ipvs_service_policy.Parse(packBytes); err != nil {
        t.Fatalf("error ipvs_service_policy.Parse: %s", err)
    } else if err := unpackService.unpack(unpackAttrs.(nlgo.AttrMap)); err != nil {
        t.Fatalf("error Service.unpack: %s", err)
    }

    testServiceEquals(t, testService, unpackService)
}

func testDestEquals (t *testing.T, testDest Dest, dest Dest) {
    if dest.Addr.String() != testDest.Addr.String() {
        t.Errorf("fail testDest.unpack(): Addr %v", dest.Addr.String())
    }
    if dest.Port != testDest.Port {
        t.Errorf("fail testDest.unpack(): Port %v", dest.Port)
    }
    if dest.FwdMethod != testDest.FwdMethod {
        t.Errorf("fail testDest.unpack(): FwdMethod %v", dest.FwdMethod)
    }
    if dest.Weight != testDest.Weight {
        t.Errorf("fail testDest.unpack(): Weight %v", dest.Weight)
    }
    if dest.UThresh != testDest.UThresh {
        t.Errorf("fail testDest.unpack(): UThresh %v", dest.UThresh)
    }
    if dest.LThresh != testDest.LThresh {
        t.Errorf("fail testDest.unpack(): LThresh %v", dest.LThresh)
    }
}

func TestDest (t *testing.T) {
    testService := Service {
        Af:     syscall.AF_INET,
    }
    testDest := Dest{
        Addr:   net.ParseIP("10.107.107.0"),
        Port:   1337,

        FwdMethod:  IP_VS_CONN_F_TUNNEL,
        Weight:     10,
        UThresh:    1000,
        LThresh:    0,
    }
    testAttrs := nlgo.AttrSlice{
        nlattr(IPVS_DEST_ATTR_ADDR, nlgo.Binary([]byte{0x0a, 0x6b, 0x6b, 0x00})),
        nlattr(IPVS_DEST_ATTR_PORT, nlgo.U16(0x3905)),
        nlattr(IPVS_DEST_ATTR_FWD_METHOD, nlgo.U32(IP_VS_CONN_F_TUNNEL)),
        nlattr(IPVS_DEST_ATTR_WEIGHT, nlgo.U32(10)),
        nlattr(IPVS_DEST_ATTR_U_THRESH, nlgo.U32(1000)),
        nlattr(IPVS_DEST_ATTR_L_THRESH, nlgo.U32(0)),
    }

    // pack
    packAttrs := testDest.attrs(&testService, true)
    packBytes := packAttrs.Bytes()

    if !bytes.Equal(packBytes, testAttrs.Bytes()) {
        t.Errorf("fail Dest.attrs(): \n%s", hex.Dump(packBytes))
    }

    // unpack
    var unpackDest Dest

    if unpackAttrs, err := ipvs_dest_policy.Parse(packBytes); err != nil {
        t.Fatalf("error ipvs_dest_policy.Parse: %s", err)
    } else if err := unpackDest.unpack(testService, unpackAttrs.(nlgo.AttrMap)); err != nil {
        t.Fatalf("error Service.unpack: %s", err)
    }

    testDestEquals(t, testDest, unpackDest)
}
<|endoftext|>


<commit_before>// Package snaker provides methods to convert CamelCase to and from snake_case.
//
// snaker takes into takes into consideration common initialisms (ie, ID, HTTP,
// ACL, etc) when converting to/from CamelCase and snake_case.
package snaker

//go:generate ./gen.sh --update

import (
	"strings"
	"unicode"
)

// CamelToSnake converts s to snake_case.
func CamelToSnake(s string) string {
	if s == "" {
		return ""
	}

	rs := []rune(s)

	var r string
	var lastWasUpper, lastWasLetter, lastWasIsm, isUpper, isLetter bool
	for i := 0; i < len(rs); {
		isUpper = unicode.IsUpper(rs[i])
		isLetter = unicode.IsLetter(rs[i])

		// append _ when last was not upper and not letter
		if (lastWasLetter && isUpper) || (lastWasIsm && isLetter) {
			r += "_"
		}

		// determine next to append to r
		var next string
		if ism := peekInitialism(rs[i:]); ism != "" && (!lastWasUpper || lastWasIsm) {
			next = ism
		} else {
			next = string(rs[i])
		}

		lastWasIsm = false
		if len(next) > 1 {
			lastWasIsm = true
		}

		lastWasUpper = isUpper
		lastWasLetter = isLetter

		r += next
		i += len(next)
	}

	return strings.ToLower(r)
}

// CamelToSnakeIdentifier converts s to its snake_case identifier.
func CamelToSnakeIdentifier(s string) string {
	return toIdentifier(CamelToSnake(s))
}

// SnakeToCamel converts s to CamelCase.
func SnakeToCamel(s string) string {
	var r string

	for _, w := range strings.Split(s, "_") {
		if w == "" {
			continue
		}

		u := strings.ToUpper(w)
		if ok := commonInitialisms[u]; ok {
			r += u
		} else {
			r += strings.ToUpper(w[:1]) + strings.ToLower(w[1:])
		}
	}

	return r
}

// SnakeToCamelIdentifier converts s to its CamelCase identifier (first
// letter is capitalized).
func SnakeToCamelIdentifier(s string) string {
	return SnakeToCamel(toIdentifier(s))
}
<commit_msg>Adding AddInitialisms func<commit_after>// Package snaker provides methods to convert CamelCase to and from snake_case.
//
// snaker takes into takes into consideration common initialisms (ie, ID, HTTP,
// ACL, etc) when converting to/from CamelCase and snake_case.
package snaker

//go:generate ./gen.sh --update

import (
	"fmt"
	"strings"
	"unicode"
)

// CamelToSnake converts s to snake_case.
func CamelToSnake(s string) string {
	if s == "" {
		return ""
	}

	rs := []rune(s)

	var r string
	var lastWasUpper, lastWasLetter, lastWasIsm, isUpper, isLetter bool
	for i := 0; i < len(rs); {
		isUpper = unicode.IsUpper(rs[i])
		isLetter = unicode.IsLetter(rs[i])

		// append _ when last was not upper and not letter
		if (lastWasLetter && isUpper) || (lastWasIsm && isLetter) {
			r += "_"
		}

		// determine next to append to r
		var next string
		if ism := peekInitialism(rs[i:]); ism != "" && (!lastWasUpper || lastWasIsm) {
			next = ism
		} else {
			next = string(rs[i])
		}

		// save for next iteration
		lastWasIsm = false
		if len(next) > 1 {
			lastWasIsm = true
		}
		lastWasUpper = isUpper
		lastWasLetter = isLetter

		r += next
		i += len(next)
	}

	return strings.ToLower(r)
}

// CamelToSnakeIdentifier converts s to its snake_case identifier.
func CamelToSnakeIdentifier(s string) string {
	return toIdentifier(CamelToSnake(s))
}

// SnakeToCamel converts s to CamelCase.
func SnakeToCamel(s string) string {
	var r string

	for _, w := range strings.Split(s, "_") {
		if w == "" {
			continue
		}

		u := strings.ToUpper(w)
		if ok := commonInitialisms[u]; ok {
			r += u
		} else {
			r += strings.ToUpper(w[:1]) + strings.ToLower(w[1:])
		}
	}

	return r
}

// SnakeToCamelIdentifier converts s to its CamelCase identifier (first
// letter is capitalized).
func SnakeToCamelIdentifier(s string) string {
	return SnakeToCamel(toIdentifier(s))
}

// AddInitialisms adds initialisms to the recognized initialisms.
func AddInitialisms(initialisms ...string) error {
	for _, s := range initialisms {
		if len(s) < minInitialismLen || len(s) > maxInitialismLen {
			return fmt.Errorf("%s does not have length between %d and %d", s, minInitialismLen, maxInitialismLen)
		}
		commonInitialisms[s] = true
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The goftp Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package server

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
	"sync"
)

// DataSocket describes a data socket is used to send non-control data between the client and
// server.
type DataSocket interface {
	Host() string

	Port() int

	// the standard io.Reader interface
	Read(p []byte) (n int, err error)

	// the standard io.Writer interface
	Write(p []byte) (n int, err error)

	// the standard io.Closer interface
	Close() error
}

type ftpActiveSocket struct {
	conn   *net.TCPConn
	host   string
	port   int
	logger Logger
}

func newActiveSocket(remote string, port int, logger Logger, sessionID string) (DataSocket, error) {
	connectTo := net.JoinHostPort(remote, strconv.Itoa(port))

	logger.Print(sessionID, "Opening active data connection to "+connectTo)

	raddr, err := net.ResolveTCPAddr("tcp", connectTo)

	if err != nil {
		logger.Print(sessionID, err)
		return nil, err
	}

	tcpConn, err := net.DialTCP("tcp", nil, raddr)

	if err != nil {
		logger.Print(sessionID, err)
		return nil, err
	}

	socket := new(ftpActiveSocket)
	socket.conn = tcpConn
	socket.host = remote
	socket.port = port
	socket.logger = logger

	return socket, nil
}

func (socket *ftpActiveSocket) Host() string {
	return socket.host
}

func (socket *ftpActiveSocket) Port() int {
	return socket.port
}

func (socket *ftpActiveSocket) Read(p []byte) (n int, err error) {
	return socket.conn.Read(p)
}

func (socket *ftpActiveSocket) Write(p []byte) (n int, err error) {
	return socket.conn.Write(p)
}

func (socket *ftpActiveSocket) Close() error {
	return socket.conn.Close()
}

type ftpPassiveSocket struct {
	conn       net.Conn
	port       int
	host       string
	ingress    chan []byte
	egress     chan []byte
	logger     Logger
	lock       sync.Mutex
	err        error
	tlsConfing *tls.Config
}

func newPassiveSocket(host string, port int, logger Logger, sessionID string, tlsConfing *tls.Config) (DataSocket, error) {
	socket := new(ftpPassiveSocket)
	socket.ingress = make(chan []byte)
	socket.egress = make(chan []byte)
	socket.logger = logger
	socket.host = host
	socket.port = port
	if err := socket.GoListenAndServe(sessionID); err != nil {
		return nil, err
	}
	return socket, nil
}

func (socket *ftpPassiveSocket) Host() string {
	return socket.host
}

func (socket *ftpPassiveSocket) Port() int {
	return socket.port
}

func (socket *ftpPassiveSocket) Read(p []byte) (n int, err error) {
	if err := socket.waitForOpenSocket(); err != nil {
		return 0, err
	}
	return socket.conn.Read(p)
}

func (socket *ftpPassiveSocket) Write(p []byte) (n int, err error) {
	if err := socket.waitForOpenSocket(); err != nil {
		return 0, err
	}
	return socket.conn.Write(p)
}

func (socket *ftpPassiveSocket) Close() error {
	if socket.conn != nil {
		return socket.conn.Close()
	}
	return nil
}

func (socket *ftpPassiveSocket) GoListenAndServe(sessionID string) (err error) {
	laddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort("", strconv.Itoa(socket.port)))
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	var listener net.Listener
	listener, err = net.ListenTCP("tcp", laddr)
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	add := listener.Addr()
	parts := strings.Split(add.String(), ":")
	port, err := strconv.Atoi(parts[len(parts)-1])
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	socket.port = port
	if socket.tlsConfing != nil {
		listener = tls.NewListener(listener, socket.tlsConfing)
	}

	go func() {
		socket.lock.Lock()
		defer socket.lock.Unlock()

		conn, err := listener.Accept()
		if err != nil {
			socket.err = err
			return
		}
		socket.err = nil
		socket.conn = conn
	}()
	return nil
}

func (socket *ftpPassiveSocket) waitForOpenSocket() error {
	socket.lock.Lock()
	defer socket.lock.Unlock()
	if socket.conn != nil {
		return nil
	}
	return socket.err
}
<commit_msg>:zap: Use sendfile syscall to send out file (#44)<commit_after>// Copyright 2018 The goftp Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package server

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
	"sync"
	"io"
)

// DataSocket describes a data socket is used to send non-control data between the client and
// server.
type DataSocket interface {
	Host() string

	Port() int

	// the standard io.Reader interface
	Read(p []byte) (n int, err error)

	// the standard io.ReaderFrom interface
	ReadFrom(r io.Reader) (int64, error)

	// the standard io.Writer interface
	Write(p []byte) (n int, err error)

	// the standard io.Closer interface
	Close() error
}

type ftpActiveSocket struct {
	conn   *net.TCPConn
	host   string
	port   int
	logger Logger
}

func newActiveSocket(remote string, port int, logger Logger, sessionID string) (DataSocket, error) {
	connectTo := net.JoinHostPort(remote, strconv.Itoa(port))

	logger.Print(sessionID, "Opening active data connection to "+connectTo)

	raddr, err := net.ResolveTCPAddr("tcp", connectTo)

	if err != nil {
		logger.Print(sessionID, err)
		return nil, err
	}

	tcpConn, err := net.DialTCP("tcp", nil, raddr)

	if err != nil {
		logger.Print(sessionID, err)
		return nil, err
	}

	socket := new(ftpActiveSocket)
	socket.conn = tcpConn
	socket.host = remote
	socket.port = port
	socket.logger = logger

	return socket, nil
}

func (socket *ftpActiveSocket) Host() string {
	return socket.host
}

func (socket *ftpActiveSocket) Port() int {
	return socket.port
}

func (socket *ftpActiveSocket) Read(p []byte) (n int, err error) {
	return socket.conn.Read(p)
}

func (socket *ftpActiveSocket) ReadFrom(r io.Reader) (int64, error) {
	return socket.conn.ReadFrom(r)
}

func (socket *ftpActiveSocket) Write(p []byte) (n int, err error) {
	return socket.conn.Write(p)
}

func (socket *ftpActiveSocket) Close() error {
	return socket.conn.Close()
}

type ftpPassiveSocket struct {
	conn       net.Conn
	port       int
	host       string
	ingress    chan []byte
	egress     chan []byte
	logger     Logger
	lock       sync.Mutex
	err        error
	tlsConfing *tls.Config
}

func newPassiveSocket(host string, port int, logger Logger, sessionID string, tlsConfing *tls.Config) (DataSocket, error) {
	socket := new(ftpPassiveSocket)
	socket.ingress = make(chan []byte)
	socket.egress = make(chan []byte)
	socket.logger = logger
	socket.host = host
	socket.port = port
	if err := socket.GoListenAndServe(sessionID); err != nil {
		return nil, err
	}
	return socket, nil
}

func (socket *ftpPassiveSocket) Host() string {
	return socket.host
}

func (socket *ftpPassiveSocket) Port() int {
	return socket.port
}

func (socket *ftpPassiveSocket) Read(p []byte) (n int, err error) {
	if err := socket.waitForOpenSocket(); err != nil {
		return 0, err
	}
	return socket.conn.Read(p)
}

func (socket *ftpPassiveSocket) ReadFrom(r io.Reader) (int64, error) {
	if err := socket.waitForOpenSocket(); err != nil {
		return 0, err
	}

	// For normal TCPConn, this will use sendfile syscall; if not,
	// it will just downgrade to normal read/write procedure
	return io.Copy(socket.conn, r)
}

func (socket *ftpPassiveSocket) Write(p []byte) (n int, err error) {
	if err := socket.waitForOpenSocket(); err != nil {
		return 0, err
	}
	return socket.conn.Write(p)
}

func (socket *ftpPassiveSocket) Close() error {
	if socket.conn != nil {
		return socket.conn.Close()
	}
	return nil
}

func (socket *ftpPassiveSocket) GoListenAndServe(sessionID string) (err error) {
	laddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort("", strconv.Itoa(socket.port)))
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	var listener net.Listener
	listener, err = net.ListenTCP("tcp", laddr)
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	add := listener.Addr()
	parts := strings.Split(add.String(), ":")
	port, err := strconv.Atoi(parts[len(parts)-1])
	if err != nil {
		socket.logger.Print(sessionID, err)
		return
	}

	socket.port = port
	if socket.tlsConfing != nil {
		listener = tls.NewListener(listener, socket.tlsConfing)
	}

	go func() {
		socket.lock.Lock()
		defer socket.lock.Unlock()

		conn, err := listener.Accept()
		if err != nil {
			socket.err = err
			return
		}
		socket.err = nil
		socket.conn = conn
	}()
	return nil
}

func (socket *ftpPassiveSocket) waitForOpenSocket() error {
	socket.lock.Lock()
	defer socket.lock.Unlock()
	if socket.conn != nil {
		return nil
	}
	return socket.err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// TODO:LIST
// parse flags
// search dirs
// search in files
// show result

// flags
var (
	root       = flag.String("root", "./", "Specify search root")
	suffix     = flag.String("filetype", "go txt", `Specify target file type into the " "`)
	suffixList []string
	gatherTarget = flag.String("key", "TODO", `Specify gather target keyword`)
)

func init() {
	var err error
	flag.Parse()
	*root, err = filepath.Abs(*root)
	if err != nil {
		log.Fatalf("init:%v\n", err)
	}
	suffixList = strings.Split(*suffix, " ")
	argsCheck()
}

// Checking after parsing flags
func argsCheck() {
	if len(flag.Args()) != 0 {
		fmt.Printf("cmd = %v\n\n", os.Args)
		fmt.Printf("-----| Unknown option |-----\n\n")
		for _, x := range flag.Args() {
			fmt.Println(x)
		}
		fmt.Printf("\n")
		fmt.Println("-----| Usage |-----")
		flag.PrintDefaults()
		os.Exit(1)
	}
}

// Use wait group!!
// もう少しシンプルにしたい
func dirsCrawl(root string) ([]string, map[string][]os.FileInfo) {

	// mux group
	var Dirlist []string
	DirsCache := make(map[string]bool)
	InfoCache := make(map[string][]os.FileInfo)
	Mux := new(sync.Mutex)

	Wg := new(sync.WaitGroup)

	var crawl func(string)
	crawl = func(dirname string) {
		defer Wg.Done()

		Mux.Lock()
		if DirsCache[dirname] {
			return
		}
		DirsCache[dirname] = true
		Mux.Unlock()

		f, err := os.Open(dirname)
		if err != nil {
			log.Printf("crawl:%v\n", err)
			return
		}
		defer func() {
			if errclose := f.Close(); errclose != nil {
				log.Println(errclose)
			}
		}()

		info, err := f.Readdir(0)
		if err != nil {
			log.Printf("crawl info:%v", err)
			return
		}
		Mux.Lock()
		InfoCache[dirname] = info
		Mux.Unlock()

		for _, x := range info {
			if x.IsDir() {
				tmp := filepath.Join(dirname, x.Name())
				Mux.Lock()
				Dirlist = append(Dirlist, tmp)
				Mux.Unlock()
				Wg.Add(1)
				go crawl(tmp)
			}
		}
	}

	Wg.Add(1)
	crawl(root)
	Wg.Wait()
	return Dirlist, InfoCache
}

// TODO:Do erase after test
func crawlshow() {
	dirslist, infomap := dirsCrawl(*root)
	fmt.Println("show directory list")
	for _, x := range dirslist {
		fmt.Println(x)
	}
	fmt.Println("root = ", *root)
	fmt.Println("length", len(dirslist))

	fmt.Println("show target files")
	var fileCount int
	for dirname, infos := range infomap {
		for _, info := range infos {
			if suffixSeacher(info.Name(), suffixList) {
				fmt.Println(filepath.Join(dirname, info.Name()))
				fileCount++
				time.Sleep(time.Millisecond)
			}
		}
	}
	fmt.Printf("all dirs = %v\nfile count = %v\n", len(dirslist), fileCount)
}

// suffixList
func suffixSeacher(filename string, targetSuffix []string) bool {
	for _, x := range targetSuffix {
		if strings.HasSuffix(filename, "."+x) {
			return true
		}
	}
	return false
}

// specify filename and target, Gather target(TODOs), return todoList.
// TODO:use goroutine
func gather(filename string, target string) ([]string, error) {
	var todoList []string
	f, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("todoGather:%v\n", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Printf("todoGather:%v", err)
		}
	}()

	for sc, i := bufio.NewScanner(f), uint(1); sc.Scan(); i++ {
		if err := sc.Err(); err != nil {
			return nil, fmt.Errorf("todoGather:%v\n", err)
		}
		if index := strings.Index(sc.Text(), target); index != -1 {
			todoList = append(todoList, fmt.Sprintf("L%v:%s", i, sc.Text()[index:]))
		}
	}
	return todoList, nil
}

// TODOGotcha!! main proc
// TODO:use goroutine
// TODO:驚くべき読みにくさ...何とかしたい
// 多分データ構造の選択を間違えてる
func mainproc() (todoMap map[string][]string, gatherErr error) {
	todoMap = make(map[string][]string)
	_, infomap := dirsCrawl(*root)
	for dirname, infos := range infomap {
		for _, info := range infos {
			if suffixSeacher(info.Name(), suffixList) {
				tmp, err := gather(filepath.Join(dirname, info.Name()), *gatherTarget)
				if err != nil {
					log.Printf("todoGather:%v", err)
					gatherErr = fmt.Errorf("todoGather:find errors. open or close or scan error.\n")
				}
				if tmp != nil {
					todoMap[filepath.Join(dirname, info.Name())] = tmp
				}
			}
		}
	}
	return todoMap, gatherErr
}
// show list!
func showTODOList(todoMap map[string][]string) {
	for filename, list := range todoMap {
		fmt.Println(filename)
		for _, s := range list {
			fmt.Println(s)
		}
		fmt.Println()
	}
	fmt.Printf("stack files=%v\n", len(todoMap))
}
func main() {
	// test show
	//crawlshow()

	todoMap, err := mainproc()
	if err != nil {
		log.Fatal(err)
	}
	showTODOList(todoMap)
}
<commit_msg>check point Changes to be committed: modified:   todogotcha.go<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// TODO:LIST
// parse flags
// search dirs
// search in files
// show result

// flags
var (
	root       = flag.String("root", "./", "Specify search root")
	suffix     = flag.String("filetype", "go txt", `Specify target file type into the " "`)
	suffixList []string
	gatherTarget = flag.String("key", "TODO", `Specify gather target keyword`)
)

func init() {
	var err error
	flag.Parse()
	*root, err = filepath.Abs(*root)
	if err != nil {
		log.Fatalf("init:%v\n", err)
	}
	suffixList = strings.Split(*suffix, " ")
	argsCheck()
}

// Checking after parsing flags
func argsCheck() {
	if len(flag.Args()) != 0 {
		fmt.Printf("cmd = %v\n\n", os.Args)
		fmt.Printf("-----| Unknown option |-----\n\n")
		for _, x := range flag.Args() {
			fmt.Println(x)
		}
		fmt.Printf("\n")
		fmt.Println("-----| Usage |-----")
		flag.PrintDefaults()
		os.Exit(1)
	}
}

// Use wait group!!
// もう少しシンプルにしたい
func dirsCrawl(root string) ([]string, map[string][]os.FileInfo) {

	// mux group
	var Dirlist []string
	DirsCache := make(map[string]bool)
	InfoCache := make(map[string][]os.FileInfo)
	Mux := new(sync.Mutex)

	Wg := new(sync.WaitGroup)

	var crawl func(string)
	crawl = func(dirname string) {
		defer Wg.Done()

		Mux.Lock()
		if DirsCache[dirname] {
			return
		}
		DirsCache[dirname] = true
		Mux.Unlock()

		f, err := os.Open(dirname)
		if err != nil {
			log.Printf("crawl:%v\n", err)
			return
		}
		defer func() {
			if errclose := f.Close(); errclose != nil {
				log.Println(errclose)
			}
		}()

		info, err := f.Readdir(0)
		if err != nil {
			log.Printf("crawl info:%v", err)
			return
		}
		Mux.Lock()
		InfoCache[dirname] = info
		Mux.Unlock()

		for _, x := range info {
			if x.IsDir() {
				tmp := filepath.Join(dirname, x.Name())
				Mux.Lock()
				Dirlist = append(Dirlist, tmp)
				Mux.Unlock()
				Wg.Add(1)
				go crawl(tmp)
			}
		}
	}

	Wg.Add(1)
	crawl(root)
	Wg.Wait()
	return Dirlist, InfoCache
}

// TODO:Do erase after test
func crawlshow() {
	dirslist, infomap := dirsCrawl(*root)
	fmt.Println("show directory list")
	for _, x := range dirslist {
		fmt.Println(x)
	}
	fmt.Println("root = ", *root)
	fmt.Println("length", len(dirslist))

	fmt.Println("show target files")
	var fileCount int
	for dirname, infos := range infomap {
		for _, info := range infos {
			if suffixSeacher(info.Name(), suffixList) {
				fmt.Println(filepath.Join(dirname, info.Name()))
				fileCount++
				time.Sleep(time.Millisecond)
			}
		}
	}
	fmt.Printf("all dirs = %v\nfile count = %v\n", len(dirslist), fileCount)
}

// suffixList
func suffixSeacher(filename string, targetSuffix []string) bool {
	for _, x := range targetSuffix {
		if strings.HasSuffix(filename, "."+x) {
			return true
		}
	}
	return false
}

// specify filename and target, Gather target(TODOs), return todoList.
// TODO:use goroutine
func gather(filename string, target string) ([]string, error) {
	var todoList []string
	f, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("todoGather:%v\n", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Printf("todoGather:%v", err)
		}
	}()

	for sc, i := bufio.NewScanner(f), uint(1); sc.Scan(); i++ {
		if err := sc.Err(); err != nil {
			return nil, fmt.Errorf("todoGather:%v\n", err)
		}
		if index := strings.Index(sc.Text(), target); index != -1 {
			todoList = append(todoList, fmt.Sprintf("L%v:%s", i, sc.Text()[index:]))
		}
	}
	return todoList, nil
}
// にゃん
func unlimitedGophersWroks(infoMap map[string][]os.FileInfo) (todoMap map[string][]string) {
	// need use goroutine
	todoMap = make(map[string][]string)
	mux := new(sync.Mutex)
	wg := new(sync.WaitGroup)

	worker := func(filepath string) {
		defer wg.Done()
		// TODO:!!
		// ファイルが多いとosのファイルディスクリプタ上限で叩かれるぅ...どうしよ...
		// スレッドを増やしまくるgoだとよくハマるっぽい
		// wgでカウント取って上限でwaitできれば良さそうだけど、カウンタは公開されてない
		// muxと自前のカウンタで何とかするか...ちょっと考える
		// 取り敢えず簡単にカウント取って一定値でwaitする?
		// os環境個別のディスクリプタ上限を取得する関数とかあれば管理できそうだけど見つかってぬぃ...
		todoList, err := gather(filepath, *gatherTarget)
		if err != nil {
			log.Println(err)
		}
		if todoList != nil {
			mux.Lock()
			todoMap[filepath] = todoList
			mux.Unlock()
		}
	}

	for dirname, infos := range infoMap {
		for _, info := range infos {
			if suffixSeacher(info.Name(), suffixList) {
				wg.Add(1)
				go worker(filepath.Join(dirname, info.Name()))
			}
		}
	}
	wg.Wait()
	return todoMap
}
func useGophersProc() (todoMap map[string][]string) {
	_, infomap := dirsCrawl(*root)
	todoMap = unlimitedGophersWroks(infomap)
	return todoMap
}

// TODOGotcha!! main proc
// TODO:use goroutine
// TODO:驚くべき読みにくさ...何とかしたい
// 多分データ構造の選択を間違えてる
func mainproc() (todoMap map[string][]string, gatherErr error) {
	todoMap = make(map[string][]string)
	_, infomap := dirsCrawl(*root)
	for dirname, infos := range infomap {
		for _, info := range infos {
			if suffixSeacher(info.Name(), suffixList) {
				tmp, err := gather(filepath.Join(dirname, info.Name()), *gatherTarget)
				if err != nil {
					log.Printf("todoGather:%v", err)
					gatherErr = fmt.Errorf("todoGather:find errors. open or close or scan error.\n")
				}
				if tmp != nil {
					todoMap[filepath.Join(dirname, info.Name())] = tmp
				}
			}
		}
	}
	return todoMap, gatherErr
}
// show list!
func showTODOList(todoMap map[string][]string) {
	for filename, list := range todoMap {
		fmt.Println(filename)
		for _, s := range list {
			fmt.Println(s)
		}
		fmt.Println()
	}
	fmt.Printf("stack files=%v\n", len(todoMap))
}
func main() {
	// test show
	//crawlshow()

//	todoMap, err := mainproc()
//	if err != nil {
//		log.Fatal(err)
//	}
	todoMap := useGophersProc()
	showTODOList(todoMap)
}
<|endoftext|>


<commit_before>package socks

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
)

const (
	socks4Version       = 4
	socks4Connect       = 1
	socks4Granted       = 90
	socks4Rejected      = 91
	socks4ConnectFailed = 92
	socks4UserIDInvalid = 93
)

var socks4Errors = []string{
	"",
	"request rejected or failed",
	"request rejected because SOCKS server cannot connect to identd on the client",
	"request rejected because the client program and identd report different user-ids",
}

// Socks4Server implements Socks4 Proxy Protocol(http://www.openssh.com/txt/socks4.protocol).
// Just support CONNECT command.
type Socks4Server struct {
	forward Dialer
}

// NewSocks4Server returns a new Socks4Server that can serve from new clients.
func NewSocks4Server(forward Dialer) (*Socks4Server, error) {
	return &Socks4Server{
		forward: forward,
	}, nil
}

// Serve with net.Listener for clients.
func (s *Socks4Server) Serve(listener net.Listener) error {
	for {
		conn, err := listener.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				continue
			} else {
				return err
			}
		}

		go serveSOCKS4Client(conn, s.forward)
	}
}

// Socks4Client implements Socks4 Proxy Protocol(http://www.openssh.com/txt/socks4.protocol).
type Socks4Client struct {
	network string
	address string
	userID  string
	forward Dialer
}

// NewSocks4Client return a new Socks4Client that implements Dialer interface.
// network must be supported by forward, address is proxy server's address, userID can empty.
func NewSocks4Client(network, address, userID string, forward Dialer) (*Socks4Client, error) {
	return &Socks4Client{
		network: network,
		address: address,
		userID:  userID,
		forward: forward,
	}, nil
}

// Dial return a new net.Conn if succeeded. network must be tcp, tcp4 or tcp6, address only is IPV4.
func (s *Socks4Client) Dial(network, address string) (net.Conn, error) {
	switch network {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, errors.New("socks: no support for SOCKS4 proxy connections of type:" + network)
	}

	host, portStr, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, errors.New("socks: failed to parse port:" + portStr)
	}
	if port < 1 || port > 0xffff {
		return nil, errors.New("socks: port number out of range:" + portStr)
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return nil, errors.New("socks: destination host invalid:" + host)
	}
	ip4 := ip.To4()
	if ip4 == nil {
		return nil, errors.New("socks:destination ip must be ipv4:" + host)
	}

	conn, err := s.forward.Dial(s.network, s.address)
	if err != nil {
		return nil, err
	}
	closeConn := &conn
	defer func() {
		if closeConn != nil {
			(*closeConn).Close()
		}
	}()

	buff := make([]byte, 0, 8+len(s.userID)+1)
	buff = append(buff, socks4Version, socks4Connect)
	buff = append(buff, byte(port>>8), byte(port))
	buff = append(buff, ip4...)
	if len(s.userID) != 0 {
		buff = append(buff, []byte(s.userID)...)
	}
	buff = append(buff, 0)

	if _, err := conn.Write(buff); err != nil {
		return nil, errors.New("socks: failed to write connect request to SOCKS4 server at: " + s.address + ": " + err.Error())
	}
	buff = buff[:8]
	if _, err := io.ReadFull(conn, buff); err != nil {
		return nil, errors.New("socks: failed to read connect reply from SOCKS4 server at: " + s.address + ": " + err.Error())
	}
	if buff[1] != socks4Granted {
		cd := int(buff[1]) - socks4Granted
		failure := "unknown error"
		if cd < len(socks4Errors) {
			failure = socks4Errors[cd]
		}
		return nil, errors.New("socks: SOCKS4 server at " + s.address + " failed to connect: " + failure)
	}

	closeConn = nil
	return conn, nil
}

func serveSOCKS4Client(conn net.Conn, forward Dialer) {
	defer conn.Close()

	reader := bufio.NewReader(conn)
	buff, err := reader.Peek(9)
	if err != nil {
		return
	}
	if buff[8] != 0 {
		if _, err = reader.ReadSlice(0); err != nil {
			return
		}
	}

	reply := make([]byte, 8)
	if buff[0] != socks4Version {
		reply[1] = socks4Rejected
		conn.Write(reply)
		return
	}
	if buff[1] != socks4Connect {
		reply[1] = socks4Rejected
		conn.Write(reply)
		return
	}

	port := uint16(buff[2])<<8 | uint16(buff[3])
	ip := buff[4:8]

	host := fmt.Sprintf("%d.%d.%d.%d:%d", ip[0], ip[1], ip[2], ip[3], port)
	dest, err := forward.Dial("tcp4", host)
	if err != nil {
		reply[1] = socks4ConnectFailed
		conn.Write(reply)
		return
	}
	defer dest.Close()

	reply[1] = socks4Granted
	if _, err = conn.Write(reply); err != nil {
		return
	}

	go func() {
		defer conn.Close()
		defer dest.Close()
		io.Copy(dest, conn)
	}()
	io.Copy(conn, dest)
}
<commit_msg>Patch for unexpected scenario<commit_after>package socks

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
)

const (
	socks4Version       = 4
	socks4Connect       = 1
	socks4Granted       = 90
	socks4Rejected      = 91
	socks4ConnectFailed = 92
	socks4UserIDInvalid = 93
)

var socks4Errors = []string{
	"",
	"request rejected or failed",
	"request rejected because SOCKS server cannot connect to identd on the client",
	"request rejected because the client program and identd report different user-ids",
}

// Socks4Server implements Socks4 Proxy Protocol(http://www.openssh.com/txt/socks4.protocol).
// Just support CONNECT command.
type Socks4Server struct {
	forward Dialer
}

// NewSocks4Server returns a new Socks4Server that can serve from new clients.
func NewSocks4Server(forward Dialer) (*Socks4Server, error) {
	return &Socks4Server{
		forward: forward,
	}, nil
}

// Serve with net.Listener for clients.
func (s *Socks4Server) Serve(listener net.Listener) error {
	for {
		conn, err := listener.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				continue
			} else {
				return err
			}
		}

		go serveSOCKS4Client(conn, s.forward)
	}
}

// Socks4Client implements Socks4 Proxy Protocol(http://www.openssh.com/txt/socks4.protocol).
type Socks4Client struct {
	network string
	address string
	userID  string
	forward Dialer
}

// NewSocks4Client return a new Socks4Client that implements Dialer interface.
// network must be supported by forward, address is proxy server's address, userID can empty.
func NewSocks4Client(network, address, userID string, forward Dialer) (*Socks4Client, error) {
	return &Socks4Client{
		network: network,
		address: address,
		userID:  userID,
		forward: forward,
	}, nil
}

// Dial return a new net.Conn if succeeded. network must be tcp, tcp4 or tcp6, address only is IPV4.
func (s *Socks4Client) Dial(network, address string) (net.Conn, error) {
	switch network {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, errors.New("socks: no support for SOCKS4 proxy connections of type:" + network)
	}

	host, portStr, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, errors.New("socks: failed to parse port:" + portStr)
	}
	if port < 1 || port > 0xffff {
		return nil, errors.New("socks: port number out of range:" + portStr)
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return nil, errors.New("socks: destination host invalid:" + host)
	}
	ip4 := ip.To4()
	if ip4 == nil {
		return nil, errors.New("socks:destination ip must be ipv4:" + host)
	}

	conn, err := s.forward.Dial(s.network, s.address)
	if err != nil {
		return nil, err
	}
	closeConn := &conn
	defer func() {
		if closeConn != nil {
			(*closeConn).Close()
		}
	}()

	buff := make([]byte, 0, 8+len(s.userID)+1)
	buff = append(buff, socks4Version, socks4Connect)
	buff = append(buff, byte(port>>8), byte(port))
	buff = append(buff, ip4...)
	if len(s.userID) != 0 {
		buff = append(buff, []byte(s.userID)...)
	}
	buff = append(buff, 0)

	if _, err := conn.Write(buff); err != nil {
		return nil, errors.New("socks: failed to write connect request to SOCKS4 server at: " + s.address + ": " + err.Error())
	}
	buff = buff[:8]
	if _, err := io.ReadFull(conn, buff); err != nil {
		return nil, errors.New("socks: failed to read connect reply from SOCKS4 server at: " + s.address + ": " + err.Error())
	}
	if buff[1] != socks4Granted {
		cd := int(buff[1]) - socks4Granted
		failure := "unknown error"
		if cd < len(socks4Errors) && cd >= 0 {
			failure = socks4Errors[cd]
		}
		return nil, errors.New("socks: SOCKS4 server at " + s.address + " failed to connect: " + failure)
	}

	closeConn = nil
	return conn, nil
}

func serveSOCKS4Client(conn net.Conn, forward Dialer) {
	defer conn.Close()

	reader := bufio.NewReader(conn)
	buff, err := reader.Peek(9)
	if err != nil {
		return
	}
	if buff[8] != 0 {
		if _, err = reader.ReadSlice(0); err != nil {
			return
		}
	}

	reply := make([]byte, 8)
	if buff[0] != socks4Version {
		reply[1] = socks4Rejected
		conn.Write(reply)
		return
	}
	if buff[1] != socks4Connect {
		reply[1] = socks4Rejected
		conn.Write(reply)
		return
	}

	port := uint16(buff[2])<<8 | uint16(buff[3])
	ip := buff[4:8]

	host := fmt.Sprintf("%d.%d.%d.%d:%d", ip[0], ip[1], ip[2], ip[3], port)
	dest, err := forward.Dial("tcp4", host)
	if err != nil {
		reply[1] = socks4ConnectFailed
		conn.Write(reply)
		return
	}
	defer dest.Close()

	reply[1] = socks4Granted
	if _, err = conn.Write(reply); err != nil {
		return
	}

	go func() {
		defer conn.Close()
		defer dest.Close()
		io.Copy(dest, conn)
	}()
	io.Copy(conn, dest)
}
<|endoftext|>


<commit_before>package monitor

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/cloudfoundry-community/go-cfclient"
)

type Report struct {
	SuccessCount int64
	FailureCount int64
	WarningCount int64
	ErrorCount   int64
	Errors       map[*Task]map[string]int // number of errors by error message by task
	Warnings     map[*Task]map[string]int // number of warnings by error message by task
	Elapsed      time.Duration
}

func (r *Report) String() string {
	total := r.SuccessCount + r.FailureCount + r.WarningCount
	s := "\nReport:\n"
	s += "==============\n"
	s += fmt.Sprintf("Total task executions: %d\n", total)
	s += fmt.Sprintf("Total successes: %d\n", r.SuccessCount)
	s += fmt.Sprintf("Total failures: %d\n", r.FailureCount)
	s += fmt.Sprintf("Total warnings: %d\n", r.WarningCount)
	s += fmt.Sprintf("Total errors: %d\n", r.ErrorCount)
	s += fmt.Sprintf("Elapsed time: %s\n", r.Elapsed.String())
	s += fmt.Sprintf("Average rate: %.2f tasks/sec\n", float64(total)/r.Elapsed.Seconds())

	if len(r.Errors) > 0 {
		s += fmt.Sprintf("\nErrors:\n")
		for task, errCounts := range r.Errors {
			s += fmt.Sprintf("\n    %s:\n", task.name)
			for err, count := range errCounts {
				s += fmt.Sprintf("        %s (%d failures)\n", err, count)
			}
		}
	}
	if len(r.Warnings) > 0 {
		s += fmt.Sprintf("\nWarnings:\n")
		for task, warningCounts := range r.Warnings {
			s += fmt.Sprintf("\n    %s:\n", task.name)
			for warning, count := range warningCounts {
				s += fmt.Sprintf("        %s (%d warnings)\n", warning, count)
			}
		}
	}
	return s
}

type TaskFunc func(cfg *cfclient.Config) error

type Task struct {
	name string
	fn   TaskFunc
}

type result struct {
	task    *Task
	err     error
	started time.Time
	ended   time.Time
}

// Monitor continuously runs each Task given to Add in parallel
// and gathers basic statistics on

type Monitor struct {
	tasks             []*Task
	taskRatePerSecond int64
	reporter          chan *Report
	queue             chan *Task
	results           chan *result
	halt              chan bool
	clientCfg         *cfclient.Config
	logger            io.Writer
	numWorkers        int
	warningMatchers   []*regexp.Regexp
}

func (m *Monitor) Add(name string, fn TaskFunc) {
	m.tasks = append(m.tasks, &Task{
		name: name,
		fn:   fn,
	})
}

func (m *Monitor) statsCollector() {
	report := &Report{
		Errors:   map[*Task]map[string]int{},
		Warnings: map[*Task]map[string]int{},
	}
	started := time.Now()
	for {
		select {
		case result := <-m.results:

			err := result.err
			if err != nil {
				msg := err.Error()

				ignored := false
				for _, warningMatcher := range m.warningMatchers {
					if warningMatcher.MatchString(msg) {
						ignored = true
						break
					}
				}
				if !ignored {
					if report.Errors[result.task] == nil {
						report.Errors[result.task] = map[string]int{}
					}
					report.Errors[result.task][msg]++
					report.ErrorCount++
					report.FailureCount++
				} else {
					if report.Warnings[result.task] == nil {
						report.Warnings[result.task] = map[string]int{}
					}
					report.Warnings[result.task][msg]++
					report.WarningCount++
				}
				fmt.Fprintf(
					m.logger,
					"%s - %s %s: %s\n",
					result.started.Format("2006-01-02 15:04:05.00 MST"),
					result.ended.Format("2006-01-02 15:04:05.00 MST"),
					result.task.name,
					msg,
				)
			} else {
				report.SuccessCount++
			}
		case <-m.halt:
			report.Elapsed = time.Since(started)
			m.reporter <- report
			return
		}
	}
}

// Run launches N worker routines to continually execute tasks and blocks until Stop is called
func (m *Monitor) Run() *Report {
	if m.halt != nil {
		panic("Run() called twice")
	}
	m.queue = make(chan *Task, m.numWorkers)
	m.halt = make(chan bool)
	for i := 0; i < m.numWorkers; i++ {
		go m.worker()
	}
	go m.statsCollector()
	go m.producer()
	return <-m.reporter
}

// worker executes a task from the queue and sends the result to the statsCollector
func (m *Monitor) worker() {
	for {
		select {
		case task := <-m.queue:
			res := &result{
				task:    task,
				started: time.Now(),
			}
			res.err = func() error {
				var cfg = *m.clientCfg
				return task.fn(&cfg)
			}()
			res.ended = time.Now()
			m.results <- res
		case <-m.halt:
			return
		}
	}
}

// producer continually fills the queue with tasks
func (m *Monitor) producer() {
	for {
		for _, task := range m.tasks {
			select {
			case <-m.halt:
				return
			case <-time.After(m.taskRateToDuration()):
				m.queue <- task
			}
		}
	}
}

func (m *Monitor) Stop() {
	select {
	case <-m.halt:
		return
	default:
		close(m.halt)
	}
}

func (m *Monitor) taskRateToDuration() time.Duration {
	return time.Second / time.Duration(m.taskRatePerSecond)
}

func NewMonitor(
	clientCfg *cfclient.Config,
	logger io.Writer,
	numWorkers int,
	warningMatchers []*regexp.Regexp,
	taskRatePerSecond int64,
) *Monitor {
	m := &Monitor{
		reporter:          make(chan *Report),
		results:           make(chan *result, numWorkers),
		clientCfg:         clientCfg,
		numWorkers:        numWorkers,
		warningMatchers:   warningMatchers,
		logger:            logger,
		taskRatePerSecond: taskRatePerSecond,
	}
	return m
}
<commit_msg>reformat availability summary<commit_after>package monitor

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/cloudfoundry-community/go-cfclient"
)

type Report struct {
	SuccessCount int64
	FailureCount int64
	WarningCount int64
	ErrorCount   int64
	Errors       map[*Task]map[string]int // number of errors by error message by task
	Warnings     map[*Task]map[string]int // number of warnings by error message by task
	Elapsed      time.Duration
}

func (r *Report) String() string {
	total := r.SuccessCount + r.FailureCount + r.WarningCount
	s := "\nReport:\n"
	s += "==============\n"
	s += fmt.Sprintf("Total task executions: %d\n", total)
	s += fmt.Sprintf("Total successes:       %d\n", r.SuccessCount)
	s += fmt.Sprintf("Total failures:        %d\n", r.FailureCount)
	s += fmt.Sprintf("Total warnings:        %d\n", r.WarningCount)
	s += fmt.Sprintf("Total errors:          %d\n", r.ErrorCount)
	s += fmt.Sprintf("Elapsed time:          %s\n", r.Elapsed.String())
	s += fmt.Sprintf("Average rate:          %.2f tasks/sec\n", float64(total)/r.Elapsed.Seconds())

	if len(r.Errors) > 0 {
		s += fmt.Sprintf("\nErrors:\n")
		for task, errCounts := range r.Errors {
			s += fmt.Sprintf("\n    %s:\n", task.name)
			for err, count := range errCounts {
				s += fmt.Sprintf("        %s (%d failures)\n", err, count)
			}
		}
	}
	if len(r.Warnings) > 0 {
		s += fmt.Sprintf("\nWarnings:\n")
		for task, warningCounts := range r.Warnings {
			s += fmt.Sprintf("\n    %s:\n", task.name)
			for warning, count := range warningCounts {
				s += fmt.Sprintf("        %s (%d warnings)\n", warning, count)
			}
		}
	}
	return s
}

type TaskFunc func(cfg *cfclient.Config) error

type Task struct {
	name string
	fn   TaskFunc
}

type result struct {
	task    *Task
	err     error
	started time.Time
	ended   time.Time
}

// Monitor continuously runs each Task given to Add in parallel
// and gathers basic statistics on

type Monitor struct {
	tasks             []*Task
	taskRatePerSecond int64
	reporter          chan *Report
	queue             chan *Task
	results           chan *result
	halt              chan bool
	clientCfg         *cfclient.Config
	logger            io.Writer
	numWorkers        int
	warningMatchers   []*regexp.Regexp
}

func (m *Monitor) Add(name string, fn TaskFunc) {
	m.tasks = append(m.tasks, &Task{
		name: name,
		fn:   fn,
	})
}

func (m *Monitor) statsCollector() {
	report := &Report{
		Errors:   map[*Task]map[string]int{},
		Warnings: map[*Task]map[string]int{},
	}
	started := time.Now()
	for {
		select {
		case result := <-m.results:

			err := result.err
			if err != nil {
				msg := err.Error()

				ignored := false
				for _, warningMatcher := range m.warningMatchers {
					if warningMatcher.MatchString(msg) {
						ignored = true
						break
					}
				}
				if !ignored {
					if report.Errors[result.task] == nil {
						report.Errors[result.task] = map[string]int{}
					}
					report.Errors[result.task][msg]++
					report.ErrorCount++
					report.FailureCount++
				} else {
					if report.Warnings[result.task] == nil {
						report.Warnings[result.task] = map[string]int{}
					}
					report.Warnings[result.task][msg]++
					report.WarningCount++
				}
				fmt.Fprintf(
					m.logger,
					"%s - %s %s: %s\n",
					result.started.Format("2006-01-02 15:04:05.00 MST"),
					result.ended.Format("2006-01-02 15:04:05.00 MST"),
					result.task.name,
					msg,
				)
			} else {
				report.SuccessCount++
			}
		case <-m.halt:
			report.Elapsed = time.Since(started)
			m.reporter <- report
			return
		}
	}
}

// Run launches N worker routines to continually execute tasks and blocks until Stop is called
func (m *Monitor) Run() *Report {
	if m.halt != nil {
		panic("Run() called twice")
	}
	m.queue = make(chan *Task, m.numWorkers)
	m.halt = make(chan bool)
	for i := 0; i < m.numWorkers; i++ {
		go m.worker()
	}
	go m.statsCollector()
	go m.producer()
	return <-m.reporter
}

// worker executes a task from the queue and sends the result to the statsCollector
func (m *Monitor) worker() {
	for {
		select {
		case task := <-m.queue:
			res := &result{
				task:    task,
				started: time.Now(),
			}
			res.err = func() error {
				var cfg = *m.clientCfg
				return task.fn(&cfg)
			}()
			res.ended = time.Now()
			m.results <- res
		case <-m.halt:
			return
		}
	}
}

// producer continually fills the queue with tasks
func (m *Monitor) producer() {
	for {
		for _, task := range m.tasks {
			select {
			case <-m.halt:
				return
			case <-time.After(m.taskRateToDuration()):
				m.queue <- task
			}
		}
	}
}

func (m *Monitor) Stop() {
	select {
	case <-m.halt:
		return
	default:
		close(m.halt)
	}
}

func (m *Monitor) taskRateToDuration() time.Duration {
	return time.Second / time.Duration(m.taskRatePerSecond)
}

func NewMonitor(
	clientCfg *cfclient.Config,
	logger io.Writer,
	numWorkers int,
	warningMatchers []*regexp.Regexp,
	taskRatePerSecond int64,
) *Monitor {
	m := &Monitor{
		reporter:          make(chan *Report),
		results:           make(chan *result, numWorkers),
		clientCfg:         clientCfg,
		numWorkers:        numWorkers,
		warningMatchers:   warningMatchers,
		logger:            logger,
		taskRatePerSecond: taskRatePerSecond,
	}
	return m
}
<|endoftext|>


<commit_before>package ssh

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/vault/helper/certutil"
	"github.com/hashicorp/vault/helper/parseutil"
	"github.com/hashicorp/vault/helper/strutil"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	"golang.org/x/crypto/ssh"
)

type creationBundle struct {
	KeyId           string
	ValidPrincipals []string
	PublicKey       ssh.PublicKey
	CertificateType uint32
	TTL             time.Duration
	Signer          ssh.Signer
	Role            *sshRole
	CriticalOptions map[string]string
	Extensions      map[string]string
}

func pathSign(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "sign/" + framework.GenericNameRegex("role"),

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.UpdateOperation: b.pathSign,
		},

		Fields: map[string]*framework.FieldSchema{
			"role": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `The desired role with configuration for this request.`,
			},
			"ttl": &framework.FieldSchema{
				Type: framework.TypeString,
				Description: `The requested Time To Live for the SSH certificate;
sets the expiration date. If not specified
the role default, backend default, or system
default TTL is used, in that order. Cannot
be later than the role max TTL.`,
			},
			"public_key": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `SSH public key that should be signed.`,
			},
			"valid_principals": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Valid principals, either usernames or hostnames, that the certificate should be signed for.`,
			},
			"cert_type": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Type of certificate to be created; either "user" or "host".`,
				Default:     "user",
			},
			"key_id": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Key id that the created certificate should have. If not specified, the display name of the token will be used.`,
			},
			"critical_options": &framework.FieldSchema{
				Type:        framework.TypeMap,
				Description: `Critical options that the certificate should be signed for.`,
			},
			"extensions": &framework.FieldSchema{
				Type:        framework.TypeMap,
				Description: `Extensions that the certificate should be signed for.`,
			},
		},

		HelpSynopsis:    `Request signing an SSH key using a certain role with the provided details.`,
		HelpDescription: `This path allows SSH keys to be signed according to the policy of the given role.`,
	}
}

func (b *backend) pathSign(req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	roleName := data.Get("role").(string)

	// Get the role
	role, err := b.getRole(req.Storage, roleName)
	if err != nil {
		return nil, err
	}
	if role == nil {
		return logical.ErrorResponse(fmt.Sprintf("Unknown role: %s", roleName)), nil
	}

	return b.pathSignCertificate(req, data, role)
}

func (b *backend) pathSignCertificate(req *logical.Request, data *framework.FieldData, role *sshRole) (*logical.Response, error) {
	publicKey := data.Get("public_key").(string)
	if publicKey == "" {
		return logical.ErrorResponse("missing public_key"), nil
	}

	userPublicKey, err := parsePublicSSHKey(publicKey)
	if err != nil {
		return logical.ErrorResponse(fmt.Sprintf("failed to parse public_key as SSH key: %s", err)), nil
	}

	// Note that these various functions always return "user errors" so we pass
	// them as 4xx values
	keyId, err := b.calculateKeyId(data, req, role, userPublicKey)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	certificateType, err := b.calculateCertificateType(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	var parsedPrincipals []string
	if certificateType == ssh.HostCert {
		parsedPrincipals, err = b.calculateValidPrincipals(data, "", role.AllowedDomains, validateValidPrincipalForHosts(role))
		if err != nil {
			return logical.ErrorResponse(err.Error()), nil
		}
	} else {
		parsedPrincipals, err = b.calculateValidPrincipals(data, role.DefaultUser, role.AllowedUsers, strutil.StrListContains)
		if err != nil {
			return logical.ErrorResponse(err.Error()), nil
		}
	}

	ttl, err := b.calculateTTL(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	criticalOptions, err := b.calculateCriticalOptions(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	extensions, err := b.calculateExtensions(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	privateKeyEntry, err := caKey(req.Storage, caPrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to read CA private key: %v", err)
	}
	if privateKeyEntry == nil || privateKeyEntry.Key == "" {
		return nil, fmt.Errorf("failed to read CA private key")
	}

	signer, err := ssh.ParsePrivateKey([]byte(privateKeyEntry.Key))
	if err != nil {
		return nil, fmt.Errorf("failed to parse stored CA private key: %v", err)
	}

	cBundle := creationBundle{
		KeyId:           keyId,
		PublicKey:       userPublicKey,
		Signer:          signer,
		ValidPrincipals: parsedPrincipals,
		TTL:             ttl,
		CertificateType: certificateType,
		Role:            role,
		CriticalOptions: criticalOptions,
		Extensions:      extensions,
	}

	certificate, err := cBundle.sign()
	if err != nil {
		return nil, err
	}

	signedSSHCertificate := ssh.MarshalAuthorizedKey(certificate)
	if len(signedSSHCertificate) == 0 {
		return nil, fmt.Errorf("error marshaling signed certificate")
	}

	response := &logical.Response{
		Data: map[string]interface{}{
			"serial_number": strconv.FormatUint(certificate.Serial, 16),
			"signed_key":    string(signedSSHCertificate),
		},
	}

	return response, nil
}

func (b *backend) calculateValidPrincipals(data *framework.FieldData, defaultPrincipal, principalsAllowedByRole string, validatePrincipal func([]string, string) bool) ([]string, error) {
	validPrincipals := ""
	validPrincipalsRaw, ok := data.GetOk("valid_principals")
	if ok {
		validPrincipals = validPrincipalsRaw.(string)
	} else {
		validPrincipals = defaultPrincipal
	}

	parsedPrincipals := strutil.ParseDedupAndSortStrings(validPrincipals, ",")
	allowedPrincipals := strutil.ParseDedupAndSortStrings(principalsAllowedByRole, ",")
	switch {
	case len(parsedPrincipals) == 0:
		// There is nothing to process
		return nil, nil
	case len(allowedPrincipals) == 0:
		// User has requested principals to be set, but role is not configured
		// with any principals
		return nil, fmt.Errorf("role is not configured to allow any principles")
	default:
		// Role was explicitly configured to allow any principal.
		if principalsAllowedByRole == "*" {
			return parsedPrincipals, nil
		}

		for _, principal := range parsedPrincipals {
			if !validatePrincipal(allowedPrincipals, principal) {
				return nil, fmt.Errorf("%v is not a valid value for valid_principals", principal)
			}
		}
		return parsedPrincipals, nil
	}
}

func validateValidPrincipalForHosts(role *sshRole) func([]string, string) bool {
	return func(allowedPrincipals []string, validPrincipal string) bool {
		for _, allowedPrincipal := range allowedPrincipals {
			if allowedPrincipal == validPrincipal && role.AllowBareDomains {
				return true
			}
			if role.AllowSubdomains && strings.HasSuffix(validPrincipal, "."+allowedPrincipal) {
				return true
			}
		}

		return false
	}
}

func (b *backend) calculateCertificateType(data *framework.FieldData, role *sshRole) (uint32, error) {
	requestedCertificateType := data.Get("cert_type").(string)

	var certificateType uint32
	switch requestedCertificateType {
	case "user":
		if !role.AllowUserCertificates {
			return 0, errors.New("cert_type 'user' is not allowed by role")
		}
		certificateType = ssh.UserCert
	case "host":
		if !role.AllowHostCertificates {
			return 0, errors.New("cert_type 'host' is not allowed by role")
		}
		certificateType = ssh.HostCert
	default:
		return 0, errors.New("cert_type must be either 'user' or 'host'")
	}

	return certificateType, nil
}

func (b *backend) calculateKeyId(data *framework.FieldData, req *logical.Request, role *sshRole, pubKey ssh.PublicKey) (string, error) {
	reqId := data.Get("key_id").(string)

	if reqId != "" {
		if !role.AllowUserKeyIDs {
			return "", fmt.Errorf("setting key_id is not allowed by role")
		}
		return reqId, nil
	}

	keyHash := sha256.Sum256(pubKey.Marshal())
	keyId := hex.EncodeToString(keyHash[:])

	if req.DisplayName != "" {
		keyId = fmt.Sprintf("%s-%s", req.DisplayName, keyId)
	}

	keyId = fmt.Sprintf("vault-%s", keyId)

	return keyId, nil
}

func (b *backend) calculateCriticalOptions(data *framework.FieldData, role *sshRole) (map[string]string, error) {
	unparsedCriticalOptions := data.Get("critical_options").(map[string]interface{})
	if len(unparsedCriticalOptions) == 0 {
		return role.DefaultCriticalOptions, nil
	}

	criticalOptions := convertMapToStringValue(unparsedCriticalOptions)

	if role.AllowedCriticalOptions != "" {
		notAllowedOptions := []string{}
		allowedCriticalOptions := strings.Split(role.AllowedCriticalOptions, ",")

		for option := range criticalOptions {
			if !strutil.StrListContains(allowedCriticalOptions, option) {
				notAllowedOptions = append(notAllowedOptions, option)
			}
		}

		if len(notAllowedOptions) != 0 {
			return nil, fmt.Errorf("Critical options not on allowed list: %v", notAllowedOptions)
		}
	}

	return criticalOptions, nil
}

func (b *backend) calculateExtensions(data *framework.FieldData, role *sshRole) (map[string]string, error) {
	unparsedExtensions := data.Get("extensions").(map[string]interface{})
	if len(unparsedExtensions) == 0 {
		return role.DefaultExtensions, nil
	}

	extensions := convertMapToStringValue(unparsedExtensions)

	if role.AllowedExtensions != "" {
		notAllowed := []string{}
		allowedExtensions := strings.Split(role.AllowedExtensions, ",")

		for extension := range extensions {
			if !strutil.StrListContains(allowedExtensions, extension) {
				notAllowed = append(notAllowed, extension)
			}
		}

		if len(notAllowed) != 0 {
			return nil, fmt.Errorf("extensions %v are not on allowed list", notAllowed)
		}
	}

	return extensions, nil
}

func (b *backend) calculateTTL(data *framework.FieldData, role *sshRole) (time.Duration, error) {

	var ttl, maxTTL time.Duration
	var ttlField string
	ttlFieldInt, ok := data.GetOk("ttl")
	if !ok {
		ttlField = role.TTL
	} else {
		ttlField = ttlFieldInt.(string)
	}

	if len(ttlField) == 0 {
		ttl = b.System().DefaultLeaseTTL()
	} else {
		var err error
		ttl, err = parseutil.ParseDurationSecond(ttlField)
		if err != nil {
			return 0, fmt.Errorf("invalid requested ttl: %s", err)
		}
	}

	if len(role.MaxTTL) == 0 {
		maxTTL = b.System().MaxLeaseTTL()
	} else {
		var err error
		maxTTL, err = parseutil.ParseDurationSecond(role.MaxTTL)
		if err != nil {
			return 0, fmt.Errorf("invalid requested max ttl: %s", err)
		}
	}

	if ttl > maxTTL {
		// Don't error if they were using system defaults, only error if
		// they specifically chose a bad TTL
		if len(ttlField) == 0 {
			ttl = maxTTL
		} else {
			return 0, fmt.Errorf("ttl is larger than maximum allowed (%d)", maxTTL/time.Second)
		}
	}

	return ttl, nil
}

func (b *creationBundle) sign() (*ssh.Certificate, error) {
	serialNumber, err := certutil.GenerateSerialNumber()
	if err != nil {
		return nil, err
	}

	now := time.Now()

	certificate := &ssh.Certificate{
		Serial:          serialNumber.Uint64(),
		Key:             b.PublicKey,
		KeyId:           b.KeyId,
		ValidPrincipals: b.ValidPrincipals,
		ValidAfter:      uint64(now.Add(-30 * time.Second).In(time.UTC).Unix()),
		ValidBefore:     uint64(now.Add(b.TTL).In(time.UTC).Unix()),
		CertType:        b.CertificateType,
		Permissions: ssh.Permissions{
			CriticalOptions: b.CriticalOptions,
			Extensions:      b.Extensions,
		},
	}

	err = certificate.SignCert(rand.Reader, b.Signer)
	if err != nil {
		return nil, fmt.Errorf("failed to generate signed SSH key")
	}

	return certificate, nil
}
<commit_msg>Parse and dedup but do not lowercase principals in SSH certs. (#2591)<commit_after>package ssh

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/vault/helper/certutil"
	"github.com/hashicorp/vault/helper/parseutil"
	"github.com/hashicorp/vault/helper/strutil"
	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	"golang.org/x/crypto/ssh"
)

type creationBundle struct {
	KeyId           string
	ValidPrincipals []string
	PublicKey       ssh.PublicKey
	CertificateType uint32
	TTL             time.Duration
	Signer          ssh.Signer
	Role            *sshRole
	CriticalOptions map[string]string
	Extensions      map[string]string
}

func pathSign(b *backend) *framework.Path {
	return &framework.Path{
		Pattern: "sign/" + framework.GenericNameRegex("role"),

		Callbacks: map[logical.Operation]framework.OperationFunc{
			logical.UpdateOperation: b.pathSign,
		},

		Fields: map[string]*framework.FieldSchema{
			"role": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `The desired role with configuration for this request.`,
			},
			"ttl": &framework.FieldSchema{
				Type: framework.TypeString,
				Description: `The requested Time To Live for the SSH certificate;
sets the expiration date. If not specified
the role default, backend default, or system
default TTL is used, in that order. Cannot
be later than the role max TTL.`,
			},
			"public_key": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `SSH public key that should be signed.`,
			},
			"valid_principals": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Valid principals, either usernames or hostnames, that the certificate should be signed for.`,
			},
			"cert_type": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Type of certificate to be created; either "user" or "host".`,
				Default:     "user",
			},
			"key_id": &framework.FieldSchema{
				Type:        framework.TypeString,
				Description: `Key id that the created certificate should have. If not specified, the display name of the token will be used.`,
			},
			"critical_options": &framework.FieldSchema{
				Type:        framework.TypeMap,
				Description: `Critical options that the certificate should be signed for.`,
			},
			"extensions": &framework.FieldSchema{
				Type:        framework.TypeMap,
				Description: `Extensions that the certificate should be signed for.`,
			},
		},

		HelpSynopsis:    `Request signing an SSH key using a certain role with the provided details.`,
		HelpDescription: `This path allows SSH keys to be signed according to the policy of the given role.`,
	}
}

func (b *backend) pathSign(req *logical.Request, data *framework.FieldData) (*logical.Response, error) {
	roleName := data.Get("role").(string)

	// Get the role
	role, err := b.getRole(req.Storage, roleName)
	if err != nil {
		return nil, err
	}
	if role == nil {
		return logical.ErrorResponse(fmt.Sprintf("Unknown role: %s", roleName)), nil
	}

	return b.pathSignCertificate(req, data, role)
}

func (b *backend) pathSignCertificate(req *logical.Request, data *framework.FieldData, role *sshRole) (*logical.Response, error) {
	publicKey := data.Get("public_key").(string)
	if publicKey == "" {
		return logical.ErrorResponse("missing public_key"), nil
	}

	userPublicKey, err := parsePublicSSHKey(publicKey)
	if err != nil {
		return logical.ErrorResponse(fmt.Sprintf("failed to parse public_key as SSH key: %s", err)), nil
	}

	// Note that these various functions always return "user errors" so we pass
	// them as 4xx values
	keyId, err := b.calculateKeyId(data, req, role, userPublicKey)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	certificateType, err := b.calculateCertificateType(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	var parsedPrincipals []string
	if certificateType == ssh.HostCert {
		parsedPrincipals, err = b.calculateValidPrincipals(data, "", role.AllowedDomains, validateValidPrincipalForHosts(role))
		if err != nil {
			return logical.ErrorResponse(err.Error()), nil
		}
	} else {
		parsedPrincipals, err = b.calculateValidPrincipals(data, role.DefaultUser, role.AllowedUsers, strutil.StrListContains)
		if err != nil {
			return logical.ErrorResponse(err.Error()), nil
		}
	}

	ttl, err := b.calculateTTL(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	criticalOptions, err := b.calculateCriticalOptions(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	extensions, err := b.calculateExtensions(data, role)
	if err != nil {
		return logical.ErrorResponse(err.Error()), nil
	}

	privateKeyEntry, err := caKey(req.Storage, caPrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to read CA private key: %v", err)
	}
	if privateKeyEntry == nil || privateKeyEntry.Key == "" {
		return nil, fmt.Errorf("failed to read CA private key")
	}

	signer, err := ssh.ParsePrivateKey([]byte(privateKeyEntry.Key))
	if err != nil {
		return nil, fmt.Errorf("failed to parse stored CA private key: %v", err)
	}

	cBundle := creationBundle{
		KeyId:           keyId,
		PublicKey:       userPublicKey,
		Signer:          signer,
		ValidPrincipals: parsedPrincipals,
		TTL:             ttl,
		CertificateType: certificateType,
		Role:            role,
		CriticalOptions: criticalOptions,
		Extensions:      extensions,
	}

	certificate, err := cBundle.sign()
	if err != nil {
		return nil, err
	}

	signedSSHCertificate := ssh.MarshalAuthorizedKey(certificate)
	if len(signedSSHCertificate) == 0 {
		return nil, fmt.Errorf("error marshaling signed certificate")
	}

	response := &logical.Response{
		Data: map[string]interface{}{
			"serial_number": strconv.FormatUint(certificate.Serial, 16),
			"signed_key":    string(signedSSHCertificate),
		},
	}

	return response, nil
}

func (b *backend) calculateValidPrincipals(data *framework.FieldData, defaultPrincipal, principalsAllowedByRole string, validatePrincipal func([]string, string) bool) ([]string, error) {
	validPrincipals := ""
	validPrincipalsRaw, ok := data.GetOk("valid_principals")
	if ok {
		validPrincipals = validPrincipalsRaw.(string)
	} else {
		validPrincipals = defaultPrincipal
	}

	parsedPrincipals := strutil.RemoveDuplicates(strutil.ParseStringSlice(validPrincipals, ","), false)
	allowedPrincipals := strutil.RemoveDuplicates(strutil.ParseStringSlice(principalsAllowedByRole, ","), false)
	switch {
	case len(parsedPrincipals) == 0:
		// There is nothing to process
		return nil, nil
	case len(allowedPrincipals) == 0:
		// User has requested principals to be set, but role is not configured
		// with any principals
		return nil, fmt.Errorf("role is not configured to allow any principles")
	default:
		// Role was explicitly configured to allow any principal.
		if principalsAllowedByRole == "*" {
			return parsedPrincipals, nil
		}

		for _, principal := range parsedPrincipals {
			if !validatePrincipal(allowedPrincipals, principal) {
				return nil, fmt.Errorf("%v is not a valid value for valid_principals", principal)
			}
		}
		return parsedPrincipals, nil
	}
}

func validateValidPrincipalForHosts(role *sshRole) func([]string, string) bool {
	return func(allowedPrincipals []string, validPrincipal string) bool {
		for _, allowedPrincipal := range allowedPrincipals {
			if allowedPrincipal == validPrincipal && role.AllowBareDomains {
				return true
			}
			if role.AllowSubdomains && strings.HasSuffix(validPrincipal, "."+allowedPrincipal) {
				return true
			}
		}

		return false
	}
}

func (b *backend) calculateCertificateType(data *framework.FieldData, role *sshRole) (uint32, error) {
	requestedCertificateType := data.Get("cert_type").(string)

	var certificateType uint32
	switch requestedCertificateType {
	case "user":
		if !role.AllowUserCertificates {
			return 0, errors.New("cert_type 'user' is not allowed by role")
		}
		certificateType = ssh.UserCert
	case "host":
		if !role.AllowHostCertificates {
			return 0, errors.New("cert_type 'host' is not allowed by role")
		}
		certificateType = ssh.HostCert
	default:
		return 0, errors.New("cert_type must be either 'user' or 'host'")
	}

	return certificateType, nil
}

func (b *backend) calculateKeyId(data *framework.FieldData, req *logical.Request, role *sshRole, pubKey ssh.PublicKey) (string, error) {
	reqId := data.Get("key_id").(string)

	if reqId != "" {
		if !role.AllowUserKeyIDs {
			return "", fmt.Errorf("setting key_id is not allowed by role")
		}
		return reqId, nil
	}

	keyHash := sha256.Sum256(pubKey.Marshal())
	keyId := hex.EncodeToString(keyHash[:])

	if req.DisplayName != "" {
		keyId = fmt.Sprintf("%s-%s", req.DisplayName, keyId)
	}

	keyId = fmt.Sprintf("vault-%s", keyId)

	return keyId, nil
}

func (b *backend) calculateCriticalOptions(data *framework.FieldData, role *sshRole) (map[string]string, error) {
	unparsedCriticalOptions := data.Get("critical_options").(map[string]interface{})
	if len(unparsedCriticalOptions) == 0 {
		return role.DefaultCriticalOptions, nil
	}

	criticalOptions := convertMapToStringValue(unparsedCriticalOptions)

	if role.AllowedCriticalOptions != "" {
		notAllowedOptions := []string{}
		allowedCriticalOptions := strings.Split(role.AllowedCriticalOptions, ",")

		for option := range criticalOptions {
			if !strutil.StrListContains(allowedCriticalOptions, option) {
				notAllowedOptions = append(notAllowedOptions, option)
			}
		}

		if len(notAllowedOptions) != 0 {
			return nil, fmt.Errorf("Critical options not on allowed list: %v", notAllowedOptions)
		}
	}

	return criticalOptions, nil
}

func (b *backend) calculateExtensions(data *framework.FieldData, role *sshRole) (map[string]string, error) {
	unparsedExtensions := data.Get("extensions").(map[string]interface{})
	if len(unparsedExtensions) == 0 {
		return role.DefaultExtensions, nil
	}

	extensions := convertMapToStringValue(unparsedExtensions)

	if role.AllowedExtensions != "" {
		notAllowed := []string{}
		allowedExtensions := strings.Split(role.AllowedExtensions, ",")

		for extension := range extensions {
			if !strutil.StrListContains(allowedExtensions, extension) {
				notAllowed = append(notAllowed, extension)
			}
		}

		if len(notAllowed) != 0 {
			return nil, fmt.Errorf("extensions %v are not on allowed list", notAllowed)
		}
	}

	return extensions, nil
}

func (b *backend) calculateTTL(data *framework.FieldData, role *sshRole) (time.Duration, error) {

	var ttl, maxTTL time.Duration
	var ttlField string
	ttlFieldInt, ok := data.GetOk("ttl")
	if !ok {
		ttlField = role.TTL
	} else {
		ttlField = ttlFieldInt.(string)
	}

	if len(ttlField) == 0 {
		ttl = b.System().DefaultLeaseTTL()
	} else {
		var err error
		ttl, err = parseutil.ParseDurationSecond(ttlField)
		if err != nil {
			return 0, fmt.Errorf("invalid requested ttl: %s", err)
		}
	}

	if len(role.MaxTTL) == 0 {
		maxTTL = b.System().MaxLeaseTTL()
	} else {
		var err error
		maxTTL, err = parseutil.ParseDurationSecond(role.MaxTTL)
		if err != nil {
			return 0, fmt.Errorf("invalid requested max ttl: %s", err)
		}
	}

	if ttl > maxTTL {
		// Don't error if they were using system defaults, only error if
		// they specifically chose a bad TTL
		if len(ttlField) == 0 {
			ttl = maxTTL
		} else {
			return 0, fmt.Errorf("ttl is larger than maximum allowed (%d)", maxTTL/time.Second)
		}
	}

	return ttl, nil
}

func (b *creationBundle) sign() (*ssh.Certificate, error) {
	serialNumber, err := certutil.GenerateSerialNumber()
	if err != nil {
		return nil, err
	}

	now := time.Now()

	certificate := &ssh.Certificate{
		Serial:          serialNumber.Uint64(),
		Key:             b.PublicKey,
		KeyId:           b.KeyId,
		ValidPrincipals: b.ValidPrincipals,
		ValidAfter:      uint64(now.Add(-30 * time.Second).In(time.UTC).Unix()),
		ValidBefore:     uint64(now.Add(b.TTL).In(time.UTC).Unix()),
		CertType:        b.CertificateType,
		Permissions: ssh.Permissions{
			CriticalOptions: b.CriticalOptions,
			Extensions:      b.Extensions,
		},
	}

	err = certificate.SignCert(rand.Reader, b.Signer)
	if err != nil {
		return nil, fmt.Errorf("failed to generate signed SSH key")
	}

	return certificate, nil
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"time"

	"k8s.io/kubernetes/pkg/api"

	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pps/persist"

	"go.pedge.io/lion/proto"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
)

var (
	// DefaultGCPolicy is the default GC policy used by a pipeline if one is not
	// specified.
	DefaultGCPolicy = &ppsclient.GCPolicy{
		// a day
		Success: &google_protobuf.Duration{
			Seconds: 24 * 60 * 60,
		},
		// a week
		Failure: &google_protobuf.Duration{
			Seconds: 7 * 24 * 60 * 60,
		},
	}
)

func (a *apiServer) runGC(ctx context.Context, pipelineInfo *ppsclient.PipelineInfo) {
	successTick := time.Tick(prototime.DurationFromProto(pipelineInfo.GcPolicy.Success))
	failureTick := time.Tick(prototime.DurationFromProto(pipelineInfo.GcPolicy.Failure))
	// wait blocks until it's time to run GC again
	wait := func() {
		select {
		case <-successTick:
		case <-failureTick:
		}
	}
	for {
		client, err := a.getPersistClient()
		if err != nil {
			protolion.Errorf("error getting persist client: %s", err)
			wait()
			continue
		}
		jobIDs, err := client.ListGCJobs(ctx, &persist.ListGCJobsRequest{
			PipelineName: pipelineInfo.Pipeline.Name,
			GcPolicy:     pipelineInfo.GcPolicy,
		})
		if err != nil {
			protolion.Errorf("error listing jobs to GC: %s", err)
			wait()
			continue
		}
		for _, jobID := range jobIDs.Jobs {
			zero := int64(0)
			falseVal := false
			go func(jobID string) {
				if err := a.kubeClient.Extensions().Jobs(a.namespace).Delete(jobID, &api.DeleteOptions{
					GracePeriodSeconds: &zero,
					OrphanDependents:   &falseVal,
				}); err != nil {
					// TODO: if the error indicates that the job has already been
					// deleted, just proceed to the next step
					protolion.Errorf("error deleting kubernetes job %s: %s", jobID, err)
					return
				}
				if _, err := client.GCJob(ctx, &ppsclient.Job{jobID}); err != nil {
					protolion.Errorf("error marking job %s as GC-ed: %s", jobID, err)
				}
				return
			}(jobID)
		}
		wait()
	}
	panic("unreachable")
}
<commit_msg>Fix job gc for a second time.<commit_after>package server

import (
	"context"
	"time"

	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pps/persist"

	"go.pedge.io/lion/proto"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
)

var (
	// DefaultGCPolicy is the default GC policy used by a pipeline if one is not
	// specified.
	DefaultGCPolicy = &ppsclient.GCPolicy{
		// a day
		Success: &google_protobuf.Duration{
			Seconds: 24 * 60 * 60,
		},
		// a week
		Failure: &google_protobuf.Duration{
			Seconds: 7 * 24 * 60 * 60,
		},
	}
)

func (a *apiServer) runGC(ctx context.Context, pipelineInfo *ppsclient.PipelineInfo) {
	successTick := time.Tick(prototime.DurationFromProto(pipelineInfo.GcPolicy.Success))
	failureTick := time.Tick(prototime.DurationFromProto(pipelineInfo.GcPolicy.Failure))
	// wait blocks until it's time to run GC again
	wait := func() {
		select {
		case <-successTick:
		case <-failureTick:
		}
	}
	for {
		client, err := a.getPersistClient()
		if err != nil {
			protolion.Errorf("error getting persist client: %s", err)
			wait()
			continue
		}
		jobIDs, err := client.ListGCJobs(ctx, &persist.ListGCJobsRequest{
			PipelineName: pipelineInfo.Pipeline.Name,
			GcPolicy:     pipelineInfo.GcPolicy,
		})
		if err != nil {
			protolion.Errorf("error listing jobs to GC: %s", err)
			wait()
			continue
		}
		for _, jobID := range jobIDs.Jobs {
			jobID := jobID
			go func() {
				jobInfo, err := client.InspectJob(ctx, &ppsclient.InspectJobRequest{
					Job: &ppsclient.Job{ID: jobID},
				})
				if err != nil {
					protolion.Errorf("error deleting job: %s", err)
					return
				}
				if err := a.deleteJob(ctx, jobInfo); err != nil {
					protolion.Errorf("error deleting job: %s", err)
					return
				}
				if _, err := client.GCJob(ctx, &ppsclient.Job{jobID}); err != nil {
					protolion.Errorf("error marking job %s as GC-ed: %s", jobID, err)
				}
				return
			}()
		}
		wait()
	}
	panic("unreachable")
}
<|endoftext|>


<commit_before>package data

import (
	"encoding/json"
)

// A Prefix defines a group of words that get automatically recognized (for example: HW, Test, Quiz)
type Prefix struct {
	ID         int      `json:"id"`
	Background string   `json:"background"`
	Color      string   `json:"color"`
	Words      []string `json:"words"`
	TimedEvent bool     `json:"timedEvent"`
	Default    bool     `json:"default"`
}

// DefaultPrefixes is the list of prefixes that all users start out with.
var DefaultPrefixes = []Prefix{
	Prefix{
		ID:         -1,
		Background: "4C6C9B",
		Color:      "FFFFFF",
		Words:      []string{"HW", "Read", "Reading"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "9ACD32",
		Color:      "FFFFFF",
		Words:      []string{"Project"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "C3A528",
		Color:      "FFFFFF",
		Words:      []string{"Report", "Essay", "Paper", "Write"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "FFA500",
		Color:      "FFFFFF",
		Words:      []string{"Quiz", "PopQuiz", "GradedHW", "GradedHomework"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "EE5D1E",
		Color:      "FFFFFF",
		Words:      []string{"Quest"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "DC143C",
		Color:      "FFFFFF",
		Words:      []string{"Test", "Final", "Exam", "Midterm"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "2AC0F1",
		Color:      "FFFFFF",
		Words:      []string{"ICA", "FieldTrip"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "2AF15E",
		Color:      "FFFFFF",
		Words:      []string{"Study", "Memorize"},
		TimedEvent: true,
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "000000",
		Color:      "00FF00",
		Words:      []string{"Trojun", "Hex"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "5000BC",
		Color:      "FFFFFF",
		Words:      []string{"OptionalHW", "Challenge"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "000099",
		Color:      "FFFFFF",
		Words:      []string{"Presentation", "Prez"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "123456",
		Color:      "FFFFFF",
		Words:      []string{"BuildSession", "Build"},
		TimedEvent: true,
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "5A1B87",
		Color:      "FFFFFF",
		Words:      []string{"Meeting", "Meet"},
		TimedEvent: true,
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "01B501",
		Color:      "FFFFFF",
		Words:      []string{"Begin", "Start", "Do"},
		TimedEvent: true,
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "E34000",
		Color:      "FFFFFF",
		Words:      []string{"Apply", "Application", "Deadline"},
		Default:    true,
	},
	Prefix{
		ID:         -1,
		Background: "3F4146",
		Color:      "FFFFFF",
		Words:      []string{"Form", "File", "Submit"},
		Default:    true,
	},
}

// FallbackBackground is the background color of a word that does not have an associated prefix.
const FallbackBackground = "FFD3BD"

// FallbackColor is the text color of a word that does not have an associated prefix.
const FallbackColor = "000000"

// GetPrefixesForUser returns a list of all prefixes for the given user, factoring in schools and custom settings.
func GetPrefixesForUser(user *User) ([]Prefix, error) {
	prefixes := DefaultPrefixes

	// check for school prefixes we want to add
	for _, school := range user.Schools {
		if school.Enabled {
			prefixes = append(prefixes, school.School.Prefixes()...)
		}
	}

	// load user settings
	rows, err := DB.Query("SELECT id, background, color, words, isTimedEvent FROM prefixes WHERE userId = ?", user.ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		resp := Prefix{}

		timedEventInt := -1
		wordsListString := ""

		rows.Scan(&resp.ID, &resp.Background, &resp.Color, &wordsListString, &timedEventInt)

		err := json.Unmarshal([]byte(wordsListString), &resp.Words)
		if err != nil {
			return nil, err
		}

		resp.TimedEvent = (timedEventInt == 1)

		prefixes = append(prefixes, resp)
	}

	return prefixes, nil
}
<commit_msg>add Upload tag, remove extra type names from tag list<commit_after>package data

import (
	"encoding/json"
)

// A Prefix defines a group of words that get automatically recognized (for example: HW, Test, Quiz)
type Prefix struct {
	ID         int      `json:"id"`
	Background string   `json:"background"`
	Color      string   `json:"color"`
	Words      []string `json:"words"`
	TimedEvent bool     `json:"timedEvent"`
	Default    bool     `json:"default"`
}

// DefaultPrefixes is the list of prefixes that all users start out with.
var DefaultPrefixes = []Prefix{
	{
		ID:         -1,
		Background: "4C6C9B",
		Color:      "FFFFFF",
		Words:      []string{"HW", "Read", "Reading"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "9ACD32",
		Color:      "FFFFFF",
		Words:      []string{"Project"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "C3A528",
		Color:      "FFFFFF",
		Words:      []string{"Report", "Essay", "Paper", "Write"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "FFA500",
		Color:      "FFFFFF",
		Words:      []string{"Quiz", "PopQuiz", "GradedHW", "GradedHomework"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "EE5D1E",
		Color:      "FFFFFF",
		Words:      []string{"Quest"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "DC143C",
		Color:      "FFFFFF",
		Words:      []string{"Test", "Final", "Exam", "Midterm"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "2AC0F1",
		Color:      "FFFFFF",
		Words:      []string{"ICA", "FieldTrip"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "2AF15E",
		Color:      "FFFFFF",
		Words:      []string{"Study", "Memorize"},
		TimedEvent: true,
		Default:    true,
	},
	{
		ID:         -1,
		Background: "000000",
		Color:      "00FF00",
		Words:      []string{"Trojun", "Hex"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "5000BC",
		Color:      "FFFFFF",
		Words:      []string{"OptionalHW", "Challenge"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "000099",
		Color:      "FFFFFF",
		Words:      []string{"Presentation", "Prez"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "123456",
		Color:      "FFFFFF",
		Words:      []string{"BuildSession", "Build"},
		TimedEvent: true,
		Default:    true,
	},
	{
		ID:         -1,
		Background: "5A1B87",
		Color:      "FFFFFF",
		Words:      []string{"Meeting", "Meet"},
		TimedEvent: true,
		Default:    true,
	},
	{
		ID:         -1,
		Background: "01B501",
		Color:      "FFFFFF",
		Words:      []string{"Begin", "Start", "Do"},
		TimedEvent: true,
		Default:    true,
	},
	{
		ID:         -1,
		Background: "E34000",
		Color:      "FFFFFF",
		Words:      []string{"Apply", "Application", "Deadline"},
		Default:    true,
	},
	{
		ID:         -1,
		Background: "3F4146",
		Color:      "FFFFFF",
		Words:      []string{"Form", "File", "Submit", "Upload"},
		Default:    true,
	},
}

// FallbackBackground is the background color of a word that does not have an associated prefix.
const FallbackBackground = "FFD3BD"

// FallbackColor is the text color of a word that does not have an associated prefix.
const FallbackColor = "000000"

// GetPrefixesForUser returns a list of all prefixes for the given user, factoring in schools and custom settings.
func GetPrefixesForUser(user *User) ([]Prefix, error) {
	prefixes := DefaultPrefixes

	// check for school prefixes we want to add
	for _, school := range user.Schools {
		if school.Enabled {
			prefixes = append(prefixes, school.School.Prefixes()...)
		}
	}

	// load user settings
	rows, err := DB.Query("SELECT id, background, color, words, isTimedEvent FROM prefixes WHERE userId = ?", user.ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		resp := Prefix{}

		timedEventInt := -1
		wordsListString := ""

		rows.Scan(&resp.ID, &resp.Background, &resp.Color, &wordsListString, &timedEventInt)

		err := json.Unmarshal([]byte(wordsListString), &resp.Words)
		if err != nil {
			return nil, err
		}

		resp.TimedEvent = (timedEventInt == 1)

		prefixes = append(prefixes, resp)
	}

	return prefixes, nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"math/rand"
	"os"

	"net/url"

	"strconv"

	"github.com/ChimeraCoder/anaconda"
	"github.com/dns-gh/tojson"
)

const (
	tweetsFilePath = "tweets.json"
)

var (
	searchTweetQueries = [...]string{
		"asteroids earth",
		"asteroids mars",
		"asteroids space",
		"asteroids nasa",
		"asteroids deadly",
	}
	maxRetweetBySearch = 3
)

// TODO factorize with load, merge and update
func loadTweets(path string) ([]anaconda.Tweet, error) {
	tweets := &[]anaconda.Tweet{}
	if _, err := os.Stat(path); os.IsNotExist(err) {
		tojson.Save(path, tweets)
	}
	err := tojson.Load(path, tweets)
	if err != nil {
		return nil, err
	}
	return *tweets, nil
}

func mergeTweets(previous, current []anaconda.Tweet) ([]anaconda.Tweet, []anaconda.Tweet) {
	merged := []anaconda.Tweet{}
	diff := []anaconda.Tweet{}
	added := map[int64]struct{}{}
	for _, v := range previous {
		added[v.Id] = struct{}{}
		merged = append(merged, v)
	}
	for _, v := range current {
		if _, ok := added[v.Id]; ok {
			continue
		}
		added[v.Id] = struct{}{}
		merged = append(merged, v)
		diff = append(diff, v)
	}
	return merged, diff
}

func updateTweets(path string, current []anaconda.Tweet) ([]anaconda.Tweet, error) {
	previous, err := loadTweets(path)
	if err != nil {
		return nil, err
	}
	merged, diff := mergeTweets(previous, current)
	tojson.Save(path, merged)
	return diff, nil
}

func getRelevantTweets() ([]anaconda.Tweet, error) {
	query := searchTweetQueries[rand.Intn(len(searchTweetQueries))]
	v := url.Values{}
	v.Set("count", strconv.Itoa(maxRetweetBySearch+2))
	results, _ := twitterAPI.GetSearch(query, v)
	return results.Statuses, nil
}

func checkRetweet() error {
	current, err := getRelevantTweets()
	if err != nil {
		return err
	}
	diff, err := updateTweets(tweetsFilePath, current)
	if err != nil {
		return err
	}
	for _, tweet := range diff {
		retweet, err := twitterAPI.Retweet(tweet.Id, false)
		if err != nil {
			log.Println("failed to retweet msg for tweet id:", tweet.Id)
		}
		log.Println("retweet (id:", tweet.Id, "):", retweet.Text)
	}
	return nil
}
<commit_msg>[srb] fix logging when error<commit_after>package main

import (
	"log"
	"math/rand"
	"os"

	"net/url"

	"strconv"

	"github.com/ChimeraCoder/anaconda"
	"github.com/dns-gh/tojson"
)

const (
	tweetsFilePath = "tweets.json"
)

var (
	searchTweetQueries = [...]string{
		"asteroids earth",
		"asteroids mars",
		"asteroids space",
		"asteroids nasa",
		"asteroids deadly",
	}
	maxRetweetBySearch = 3
)

// TODO factorize with load, merge and update
func loadTweets(path string) ([]anaconda.Tweet, error) {
	tweets := &[]anaconda.Tweet{}
	if _, err := os.Stat(path); os.IsNotExist(err) {
		tojson.Save(path, tweets)
	}
	err := tojson.Load(path, tweets)
	if err != nil {
		return nil, err
	}
	return *tweets, nil
}

func mergeTweets(previous, current []anaconda.Tweet) ([]anaconda.Tweet, []anaconda.Tweet) {
	merged := []anaconda.Tweet{}
	diff := []anaconda.Tweet{}
	added := map[int64]struct{}{}
	for _, v := range previous {
		added[v.Id] = struct{}{}
		merged = append(merged, v)
	}
	for _, v := range current {
		if _, ok := added[v.Id]; ok {
			continue
		}
		added[v.Id] = struct{}{}
		merged = append(merged, v)
		diff = append(diff, v)
	}
	return merged, diff
}

func updateTweets(path string, current []anaconda.Tweet) ([]anaconda.Tweet, error) {
	previous, err := loadTweets(path)
	if err != nil {
		return nil, err
	}
	merged, diff := mergeTweets(previous, current)
	tojson.Save(path, merged)
	return diff, nil
}

func getRelevantTweets() ([]anaconda.Tweet, error) {
	query := searchTweetQueries[rand.Intn(len(searchTweetQueries))]
	v := url.Values{}
	v.Set("count", strconv.Itoa(maxRetweetBySearch+2))
	results, _ := twitterAPI.GetSearch(query, v)
	return results.Statuses, nil
}

func checkRetweet() error {
	current, err := getRelevantTweets()
	if err != nil {
		return err
	}
	diff, err := updateTweets(tweetsFilePath, current)
	if err != nil {
		return err
	}
	for _, tweet := range diff {
		retweet, err := twitterAPI.Retweet(tweet.Id, false)
		if err != nil {
			log.Printf("failed to retweet msg for tweet (id:%d), error: %v\n", tweet.Id, err)
			continue
		}
		log.Println("retweet (id:", tweet.Id, "):", retweet.Text)
	}
	return nil
}
<|endoftext|>