<commit_before>// Google Cloud Messaging for application servers implemented using the
// Go programming language.
package gcm

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"
)

const (
	// GcmSendEndpoint is the endpoint for sending messages to the GCM server.
	GcmSendEndpoint = "https://android.googleapis.com/gcm/send"
	// Initial delay before first retry, without jitter.
	backoffInitialDelay = 1000
	// Maximum delay before a retry.
	maxBackoffDelay = 1024000
)

// Declared as a mutable variable for testing purposes.
var gcmSendEndpoint = GcmSendEndpoint

// Sender abstracts the interaction between the application server and the
// GCM server. The developer must obtain an API key from the Google APIs
// Console page and pass it to the Sender so that it can perform authorized
// requests on the application server's behalf. To send a message to one or
// more devices use the Sender's Send or SendNoRetry methods.
//
// If the Http field is nil, a zeroed http.Client will be allocated and used
// to send messages. If your application server runs on Google AppEngine,
// you must use the "appengine/urlfetch" package to create the *http.Client
// as follows:
//
//	func handler(w http.ResponseWriter, r *http.Request) {
//		c := appengine.NewContext(r)
//		client := urlfetch.Client(c)
//		sender := &gcm.Sender{ApiKey: key, Http: client}
//
//		/* ... */
//	}
type Sender struct {
	ApiKey string
	Http   *http.Client
}

// SendNoRetry sends a message to the GCM server without retrying in case of
// service unavailability. A non-nil error is returned if a non-recoverable
// error occurs (i.e. if the response status is not "200 OK").
func (s *Sender) SendNoRetry(msg *Message) (*Response, error) {
	if err := checkSender(s); err != nil {
		return nil, err
	} else if err := checkMessage(msg); err != nil {
		return nil, err
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", gcmSendEndpoint, bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", fmt.Sprintf("key=%s", s.ApiKey))
	req.Header.Add("Content-Type", "application/json")

	resp, err := s.Http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%d error: %s", resp.StatusCode, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	response := new(Response)
	err = json.Unmarshal(body, response)
	return response, err
}

// Send sends a message to the GCM server, retrying in case of service
// unavailability. A non-nil error is returned if a non-recoverable
// error occurs (i.e. if the response status is not "200 OK").
//
// Note that messages are retried using exponential backoff, and as a
// result, this method may block for several seconds.
func (s *Sender) Send(msg *Message, retries int) (*Response, error) {
	if err := checkSender(s); err != nil {
		return nil, err
	} else if err := checkMessage(msg); err != nil {
		return nil, err
	} else if retries < 0 {
		return nil, errors.New("'retries' must not be negative.")
	}

	// Send the message for the first time.
	resp, err := s.SendNoRetry(msg)
	if err != nil {
		return nil, err
	} else if resp.Failure == 0 || retries == 0 {
		return resp, nil
	}

	// One or more messages failed to send.
	regIDs := msg.RegistrationIDs
	allResults := make(map[string]Result, len(regIDs))
	backoff := backoffInitialDelay
	for i := 0; updateStatus(msg, resp, allResults) > 0 && i < retries; i++ {
		sleepTime := backoff/2 + rand.Intn(backoff)
		time.Sleep(time.Duration(sleepTime) * time.Millisecond)
		backoff = min(2*backoff, maxBackoffDelay)
		if resp, err = s.SendNoRetry(msg); err != nil {
			msg.RegistrationIDs = regIDs
			return nil, err
		}
	}

	// Bring the message back to its original state.
	msg.RegistrationIDs = regIDs

	// Create a Response containing the overall results.
	finalResults := make([]Result, len(regIDs))
	var success, failure, canonicalIDs int
	for i := 0; i < len(regIDs); i++ {
		result, _ := allResults[regIDs[i]]
		finalResults[i] = result
		if result.MessageID != "" {
			if result.RegistrationID != "" {
				canonicalIDs++
			}
			success++
		} else {
			failure++
		}
	}

	return &Response{
		// Return the most recent multicast id.
		MulticastID:  resp.MulticastID,
		Success:      success,
		Failure:      failure,
		CanonicalIDs: canonicalIDs,
		Results:      finalResults,
	}, nil
}

// updateStatus updates the status of the messages sent to devices and
// returns the number of recoverable errors that could be retried.
func updateStatus(msg *Message, resp *Response, allResults map[string]Result) int {
	unsentRegIDs := make([]string, 0, resp.Failure)
	for i := 0; i < len(resp.Results); i++ {
		regID := msg.RegistrationIDs[i]
		allResults[regID] = resp.Results[i]
		if resp.Results[i].Error == "Unavailable" {
			unsentRegIDs = append(unsentRegIDs, regID)
		}
	}
	msg.RegistrationIDs = unsentRegIDs
	return len(unsentRegIDs)
}

// min returns the smaller of two integers. For exciting religious wars
// about why this wasn't included in the "math" package, see this thread:
// https://groups.google.com/d/topic/golang-nuts/dbyqx_LGUxM/discussion
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// checkSender returns an error if the sender is not well-formed and
// initializes a zeroed http.Client if one has not been provided.
func checkSender(sender *Sender) error {
	if sender.ApiKey == "" {
		return errors.New("the sender's API key must not be empty")
	}
	if sender.Http == nil {
		sender.Http = new(http.Client)
	}
	return nil
}

// checkMessage returns an error if the message is not well-formed.
func checkMessage(msg *Message) error {
	if msg == nil {
		return errors.New("the message must not be nil")
	} else if msg.RegistrationIDs == nil {
		return errors.New("the message's RegistrationIDs field must not be nil")
	} else if len(msg.RegistrationIDs) == 0 {
		return errors.New("the message must specify at least one registration ID")
	} else if len(msg.RegistrationIDs) > 1000 {
		return errors.New("the message may specify at most 1000 registration IDs")
	} else if msg.TimeToLive < 0 || 2419200 < msg.TimeToLive {
		return errors.New("the message's TimeToLive field must be an integer " +
			"between 0 and 2419200 (4 weeks)")
	}
	return nil
}
<commit_msg>Add typed errors for known error types<commit_after>// Google Cloud Messaging for application servers implemented using the
// Go programming language.
package gcm

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"
)

const (
	// GcmSendEndpoint is the endpoint for sending messages to the GCM server.
	GcmSendEndpoint = "https://android.googleapis.com/gcm/send"
	// Initial delay before first retry, without jitter.
	backoffInitialDelay = 1000
	// Maximum delay before a retry.
	maxBackoffDelay = 1024000
)

// Declared as a mutable variable for testing purposes.
var gcmSendEndpoint = GcmSendEndpoint

// Errors
type JSONParseError struct{ error }
type UnauthorizedError struct{ error }
type UnknownError struct{ error }

const (
	ResponseErrorMissingRegistration = "MissingRegistration"
	ResponseErrorInvalidRegistration = "InvalidRegistration"
	ResponseErrorMismatchSenderID    = "MismatchSenderId"
	ResponseErrorNotRegistered       = "NotRegistered"
	ResponseErrorMessageTooBig       = "MessageTooBig"
	ResponseErrorInvalidDataKey      = "InvalidDataKey"
	ResponseErrorInvalidTTL          = "InvalidTtl"
	ResponseErrorUnavailable         = "Unavailable"
	ResponseErrorInternalServerError = "InternalServerError"
	ResponseErrorInvalidPackageName  = "InvalidPackageName"
)

// Sender abstracts the interaction between the application server and the
// GCM server. The developer must obtain an API key from the Google APIs
// Console page and pass it to the Sender so that it can perform authorized
// requests on the application server's behalf. To send a message to one or
// more devices use the Sender's Send or SendNoRetry methods.
//
// If the Http field is nil, a zeroed http.Client will be allocated and used
// to send messages. If your application server runs on Google AppEngine,
// you must use the "appengine/urlfetch" package to create the *http.Client
// as follows:
//
//	func handler(w http.ResponseWriter, r *http.Request) {
//		c := appengine.NewContext(r)
//		client := urlfetch.Client(c)
//		sender := &gcm.Sender{ApiKey: key, Http: client}
//
//		/* ... */
//	}
type Sender struct {
	ApiKey string
	Http   *http.Client
}

// SendNoRetry sends a message to the GCM server without retrying in case of
// service unavailability. A non-nil error is returned if a non-recoverable
// error occurs (i.e. if the response status is not "200 OK").
func (s *Sender) SendNoRetry(msg *Message) (*Response, error) {
	if err := checkSender(s); err != nil {
		return nil, err
	} else if err := checkMessage(msg); err != nil {
		return nil, err
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", gcmSendEndpoint, bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", fmt.Sprintf("key=%s", s.ApiKey))
	req.Header.Add("Content-Type", "application/json")

	resp, err := s.Http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	switch resp.StatusCode {
	case http.StatusBadRequest:
		return nil, JSONParseError{fmt.Errorf("%d error: %s", resp.StatusCode, resp.Status)}
	case http.StatusUnauthorized:
		return nil, UnauthorizedError{fmt.Errorf("%d error: %s", resp.StatusCode, resp.Status)}
	case http.StatusOK:
	default:
		return nil, UnknownError{fmt.Errorf("%d error: %s", resp.StatusCode, resp.Status)}
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	response := new(Response)
	err = json.Unmarshal(body, response)
	return response, err
}

// Send sends a message to the GCM server, retrying in case of service
// unavailability. A non-nil error is returned if a non-recoverable
// error occurs (i.e. if the response status is not "200 OK").
//
// Note that messages are retried using exponential backoff, and as a
// result, this method may block for several seconds.
func (s *Sender) Send(msg *Message, retries int) (*Response, error) {
	if err := checkSender(s); err != nil {
		return nil, err
	} else if err := checkMessage(msg); err != nil {
		return nil, err
	} else if retries < 0 {
		return nil, errors.New("'retries' must not be negative.")
	}

	// Send the message for the first time.
	resp, err := s.SendNoRetry(msg)
	if err != nil {
		return nil, err
	} else if resp.Failure == 0 || retries == 0 {
		return resp, nil
	}

	// One or more messages failed to send.
	regIDs := msg.RegistrationIDs
	allResults := make(map[string]Result, len(regIDs))
	backoff := backoffInitialDelay
	for i := 0; updateStatus(msg, resp, allResults) > 0 && i < retries; i++ {
		sleepTime := backoff/2 + rand.Intn(backoff)
		time.Sleep(time.Duration(sleepTime) * time.Millisecond)
		backoff = min(2*backoff, maxBackoffDelay)
		if resp, err = s.SendNoRetry(msg); err != nil {
			msg.RegistrationIDs = regIDs
			return nil, err
		}
	}

	// Bring the message back to its original state.
	msg.RegistrationIDs = regIDs

	// Create a Response containing the overall results.
	finalResults := make([]Result, len(regIDs))
	var success, failure, canonicalIDs int
	for i := 0; i < len(regIDs); i++ {
		result, _ := allResults[regIDs[i]]
		finalResults[i] = result
		if result.MessageID != "" {
			if result.RegistrationID != "" {
				canonicalIDs++
			}
			success++
		} else {
			failure++
		}
	}

	return &Response{
		// Return the most recent multicast id.
		MulticastID:  resp.MulticastID,
		Success:      success,
		Failure:      failure,
		CanonicalIDs: canonicalIDs,
		Results:      finalResults,
	}, nil
}

// updateStatus updates the status of the messages sent to devices and
// returns the number of recoverable errors that could be retried.
func updateStatus(msg *Message, resp *Response, allResults map[string]Result) int {
	unsentRegIDs := make([]string, 0, resp.Failure)
	for i := 0; i < len(resp.Results); i++ {
		regID := msg.RegistrationIDs[i]
		allResults[regID] = resp.Results[i]
		if resp.Results[i].Error == "Unavailable" {
			unsentRegIDs = append(unsentRegIDs, regID)
		}
	}
	msg.RegistrationIDs = unsentRegIDs
	return len(unsentRegIDs)
}

// min returns the smaller of two integers. For exciting religious wars
// about why this wasn't included in the "math" package, see this thread:
// https://groups.google.com/d/topic/golang-nuts/dbyqx_LGUxM/discussion
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// checkSender returns an error if the sender is not well-formed and
// initializes a zeroed http.Client if one has not been provided.
func checkSender(sender *Sender) error {
	if sender.ApiKey == "" {
		return errors.New("the sender's API key must not be empty")
	}
	if sender.Http == nil {
		sender.Http = new(http.Client)
	}
	return nil
}

// checkMessage returns an error if the message is not well-formed.
func checkMessage(msg *Message) error {
	if msg == nil {
		return errors.New("the message must not be nil")
	} else if msg.RegistrationIDs == nil {
		return errors.New("the message's RegistrationIDs field must not be nil")
	} else if len(msg.RegistrationIDs) == 0 {
		return errors.New("the message must specify at least one registration ID")
	} else if len(msg.RegistrationIDs) > 1000 {
		return errors.New("the message may specify at most 1000 registration IDs")
	} else if msg.TimeToLive < 0 || 2419200 < msg.TimeToLive {
		return errors.New("the message's TimeToLive field must be an integer " +
			"between 0 and 2419200 (4 weeks)")
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS server implementation.

package dns

import (
	"github.com/miekg/radix"
	"io"
	"net"
	"time"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the net.Addr of the client that sent the current request.
	RemoteAddr() net.Addr
	// TsigStatus returns the status of the Tsig.
	TsigStatus() error
	// Write writes a reply back to the client.
	Write(*Msg) error
	// WriteBuf writes a raw buffer back to the client.
	WriteBuf([]byte) error
}

type conn struct {
	remoteAddr net.Addr          // address of the client
	handler    Handler           // request handler
	request    []byte            // bytes read
	_UDP       *net.UDPConn      // i/o connection if UDP was used
	_TCP       *net.TCPConn      // i/o connection if TCP was used
	hijacked   bool              // connection has been hijacked by hander TODO(mg)
	tsigSecret map[string]string // the tsig secrets
}

type response struct {
	conn           *conn
	req            *Msg
	tsigStatus     error
	tsigTimersOnly bool
	tsigRequestMAC string
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name. ServeMux is DNSSEC aware, meaning
// that queries for the DS record are redirected to the parent zone (if that
// is also registered), otherwise the child gets the query.
type ServeMux struct {
	m *radix.Radix
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{m: radix.New()} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, r)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Failed is a helper handler that returns an answer with
// RCODE = servfail for every request.
func Failed(w ResponseWriter, r *Msg) {
	m := new(Msg)
	m.SetRcode(r, RcodeServerFailure)
	w.Write(m)
}

// FailedHandler returns HandlerFunc with Failed.
func FailedHandler() Handler { return HandlerFunc(Failed) }

// Start a server on addresss and network speficied. Invoke handler
// for any incoming queries.
func ListenAndServe(addr string, network string, handler Handler) error {
	server := &Server{Addr: addr, Net: network, Handler: handler}
	return server.ListenAndServe()
}

func (mux *ServeMux) match(zone string, t uint16) Handler {
	// Exact match
	zone = toRadixName(zone)
	if h := mux.m.Find(zone); h != nil && h.Value != nil {
		// If we got queried for a DS record, we must see if we
		// if we also serve the parent. We then redirect it.
		if t == TypeDS {
			if d := h.Up(); d != nil {
				return d.Value.(Handler)
			}
		}
		return h.Value.(Handler)
	}
	// Best matching one.
	if h := mux.m.Predecessor(zone); h != nil && h.Value != nil {
		return h.Value.(Handler)
	}
	return nil
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m.Insert(toRadixName(Fqdn(pattern)), handler)
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// HandleRemove deregistrars the handler specific for pattern from the ServeMux.
func (mux *ServeMux) HandleRemove(pattern string) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	// if its there, its gone
	mux.m.Remove(toRadixName(Fqdn(pattern)))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message. For DS queries
// a parent zone is sought.
// If no handler is found a standard SERVFAIL message is returned
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name, request.Question[0].Qtype)
	if h == nil {
		h = FailedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle registers the handler with the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patterns are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

// HandleRemove deregisters the handle with the given pattern
// in the DefaultServeMux.
func HandleRemove(pattern string) { DefaultServeMux.HandleRemove(pattern) }

// HandleFunc registers the handler function with the given pattern
// in the DefaultServeMux.
func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// A Server defines parameters for running an DNS server.
type Server struct {
	Addr         string            // address to listen on, ":dns" if empty
	Net          string            // if "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler           // handler to invoke, dns.DefaultServeMux if nil
	UDPSize      int               // default buffer size to use to read incoming UDP messages
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

// ListenAndServe starts a nameserver on the configured address in *Server.
func (srv *Server) ListenAndServe() error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Net {
	case "tcp", "tcp4", "tcp6":
		a, e := net.ResolveTCPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp", "udp4", "udp6":
		a, e := net.ResolveUDPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return &Error{Err: "bad network"}
}

// ServeTCP starts a TCP listener for the server.
// Each request is handled in a seperate goroutine.
func (srv *Server) ServeTCP(l *net.TCPListener) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil || n == 0 {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

// ServeUDP starts a UDP listener for the server.
// Each request is handled in a seperate goroutine.
func (srv *Server) ServeUDP(l *net.UDPConn) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	if srv.UDPSize == 0 {
		srv.UDPSize = udpMsgSize
	}
	for {
		m := make([]byte, srv.UDPSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil || n == 0 {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		d, err := newConn(nil, l, a, m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler, tsig map[string]string) (*conn, error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	c.tsigSecret = tsig
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
	// for block to make it easy to break out to close the tcp connection
	for {
		// Request has been read in ServeUDP or ServeTCP
		w := new(response)
		w.conn = c
		req := new(Msg)
		if !req.Unpack(c.request) {
			// Send a format error back
			x := new(Msg)
			x.SetRcodeFormatError(req)
			w.Write(x)
			break
		}

		w.tsigStatus = nil
		if t := req.IsTsig(); t != nil {
			secret := t.Hdr.Name
			if _, ok := w.conn.tsigSecret[secret]; !ok {
				w.tsigStatus = ErrKeyAlg
			}
			w.tsigStatus = TsigVerify(c.request, w.conn.tsigSecret[secret], "", false)
			w.tsigTimersOnly = false // Will this ever be true?
			w.tsigRequestMAC = req.Extra[len(req.Extra)-1].(*RR_TSIG).MAC
		}
		w.req = req
		c.handler.ServeDNS(w, w.req) // this does the writing back to the client
		if c.hijacked {
			return
		}
		break
	}
	if c._TCP != nil {
		c.close() // Listen and Serve is closed then
	}
}

// Write implements the ResponseWriter.Write method.
func (w *response) Write(m *Msg) (err error) {
	var (
		data []byte
		ok   bool
	)
	if m == nil {
		return &Error{Err: "nil message"}
	}
	if t := m.IsTsig(); t != nil {
		data, w.tsigRequestMAC, err = TsigGenerate(m, w.conn.tsigSecret[t.Hdr.Name], w.tsigRequestMAC, w.tsigTimersOnly)
		if err != nil {
			return err
		}
	} else {
		data, ok = m.Pack()
		if !ok {
			return ErrPack
		}
	}
	return w.WriteBuf(data)
}

// WriteBuf implements the ResponseWriter.WriteBuf method.
func (w *response) WriteBuf(m []byte) (err error) {
	if m == nil {
		return &Error{Err: "nil message"}
	}
	switch {
	case w.conn._UDP != nil:
		_, err := w.conn._UDP.WriteTo(m, w.conn.remoteAddr)
		if err != nil {
			return err
		}
	case w.conn._TCP != nil:
		if len(m) > MaxMsgSize {
			return ErrBuf
		}
		l := make([]byte, 2)
		l[0], l[1] = packUint16(uint16(len(m)))
		n, err := w.conn._TCP.Write(l)
		if err != nil {
			return err
		}
		if n != 2 {
			return io.ErrShortWrite
		}
		n, err = w.conn._TCP.Write(m)
		if err != nil {
			return err
		}
		i := n
		if i < len(m) {
			j, err := w.conn._TCP.Write(m[i:len(m)])
			if err != nil {
				return err
			}
			i += j
		}
		n = i
	}
	return nil
}

// RemoteAddr implements the ResponseWriter.RemoteAddr method.
func (w *response) RemoteAddr() net.Addr { return w.conn.remoteAddr }

// TsigStatus implements the ResponseWriter.TsigStatus method.
func (w *response) TsigStatus() error { return w.tsigStatus }
<commit_msg>Add TsigTimersOnly function<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS server implementation.

package dns

import (
	"github.com/miekg/radix"
	"io"
	"net"
	"time"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the net.Addr of the client that sent the current request.
	RemoteAddr() net.Addr
	// Write writes a reply back to the client.
	Write(*Msg) error
	// WriteBuf writes a raw buffer back to the client.
	WriteBuf([]byte) error
	// TsigStatus returns the status of the Tsig.
	TsigStatus() error
	// TsigTimersOnly sets the tsig timers only boolean.
	TsigTimersOnly(bool)
}

type conn struct {
	remoteAddr net.Addr          // address of the client
	handler    Handler           // request handler
	request    []byte            // bytes read
	_UDP       *net.UDPConn      // i/o connection if UDP was used
	_TCP       *net.TCPConn      // i/o connection if TCP was used
	hijacked   bool              // connection has been hijacked by hander TODO(mg)
	tsigSecret map[string]string // the tsig secrets
}

type response struct {
	conn           *conn
	req            *Msg
	tsigStatus     error
	tsigTimersOnly bool
	tsigRequestMAC string
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name. ServeMux is DNSSEC aware, meaning
// that queries for the DS record are redirected to the parent zone (if that
// is also registered), otherwise the child gets the query.
type ServeMux struct {
	m *radix.Radix
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{m: radix.New()} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, r)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Failed is a helper handler that returns an answer with
// RCODE = servfail for every request.
func Failed(w ResponseWriter, r *Msg) {
	m := new(Msg)
	m.SetRcode(r, RcodeServerFailure)
	w.Write(m)
}

// FailedHandler returns HandlerFunc with Failed.
func FailedHandler() Handler { return HandlerFunc(Failed) }

// Start a server on addresss and network speficied. Invoke handler
// for any incoming queries.
func ListenAndServe(addr string, network string, handler Handler) error {
	server := &Server{Addr: addr, Net: network, Handler: handler}
	return server.ListenAndServe()
}

func (mux *ServeMux) match(zone string, t uint16) Handler {
	// Exact match
	zone = toRadixName(zone)
	if h := mux.m.Find(zone); h != nil && h.Value != nil {
		// If we got queried for a DS record, we must see if we
		// if we also serve the parent. We then redirect it.
		if t == TypeDS {
			if d := h.Up(); d != nil {
				return d.Value.(Handler)
			}
		}
		return h.Value.(Handler)
	}
	// Best matching one.
	if h := mux.m.Predecessor(zone); h != nil && h.Value != nil {
		return h.Value.(Handler)
	}
	return nil
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m.Insert(toRadixName(Fqdn(pattern)), handler)
}

// Handle adds a handler to the ServeMux for pattern.
func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// HandleRemove deregistrars the handler specific for pattern from the ServeMux.
func (mux *ServeMux) HandleRemove(pattern string) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	// if its there, its gone
	mux.m.Remove(toRadixName(Fqdn(pattern)))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message. For DS queries
// a parent zone is sought.
// If no handler is found a standard SERVFAIL message is returned
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name, request.Question[0].Qtype)
	if h == nil {
		h = FailedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle registers the handler with the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patterns are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

// HandleRemove deregisters the handle with the given pattern
// in the DefaultServeMux.
func HandleRemove(pattern string) { DefaultServeMux.HandleRemove(pattern) }

// HandleFunc registers the handler function with the given pattern
// in the DefaultServeMux.
func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// A Server defines parameters for running an DNS server.
type Server struct {
	Addr         string            // address to listen on, ":dns" if empty
	Net          string            // if "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler           // handler to invoke, dns.DefaultServeMux if nil
	UDPSize      int               // default buffer size to use to read incoming UDP messages
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

// ListenAndServe starts a nameserver on the configured address in *Server.
func (srv *Server) ListenAndServe() error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Net {
	case "tcp", "tcp4", "tcp6":
		a, e := net.ResolveTCPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp", "udp4", "udp6":
		a, e := net.ResolveUDPAddr(srv.Net, addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP(srv.Net, a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return &Error{Err: "bad network"}
}

// ServeTCP starts a TCP listener for the server.
// Each request is handled in a seperate goroutine.
func (srv *Server) ServeTCP(l *net.TCPListener) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil || n == 0 {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

// ServeUDP starts a UDP listener for the server.
// Each request is handled in a seperate goroutine.
func (srv *Server) ServeUDP(l *net.UDPConn) error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	if srv.UDPSize == 0 {
		srv.UDPSize = udpMsgSize
	}
	for {
		m := make([]byte, srv.UDPSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil || n == 0 {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadDeadline(time.Now().Add(srv.ReadTimeout))
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteDeadline(time.Now().Add(srv.WriteTimeout))
		}
		d, err := newConn(nil, l, a, m, handler, srv.TsigSecret)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler, tsig map[string]string) (*conn, error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	c.tsigSecret = tsig
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
	// for block to make it easy to break out to close the tcp connection
	for {
		// Request has been read in ServeUDP or ServeTCP
		w := new(response)
		w.conn = c
		req := new(Msg)
		if !req.Unpack(c.request) {
			// Send a format error back
			x := new(Msg)
			x.SetRcodeFormatError(req)
			w.Write(x)
			break
		}

		w.tsigStatus = nil
		if t := req.IsTsig(); t != nil {
			secret := t.Hdr.Name
			if _, ok := w.conn.tsigSecret[secret]; !ok {
				w.tsigStatus = ErrKeyAlg
			}
			w.tsigStatus = TsigVerify(c.request, w.conn.tsigSecret[secret], "", false)
			w.tsigTimersOnly = false
			w.tsigRequestMAC = req.Extra[len(req.Extra)-1].(*RR_TSIG).MAC
		}
		w.req = req
		c.handler.ServeDNS(w, w.req) // this does the writing back to the client
		if c.hijacked {
			return
		}
		break
	}
	if c._TCP != nil {
		c.close() // Listen and Serve is closed then
	}
}

// Write implements the ResponseWriter.Write method.
func (w *response) Write(m *Msg) (err error) {
	var (
		data []byte
		ok   bool
	)
	if m == nil {
		return &Error{Err: "nil message"}
	}
	if t := m.IsTsig(); t != nil {
		data, w.tsigRequestMAC, err = TsigGenerate(m, w.conn.tsigSecret[t.Hdr.Name], w.tsigRequestMAC, w.tsigTimersOnly)
		if err != nil {
			return err
		}
	} else {
		data, ok = m.Pack()
		if !ok {
			return ErrPack
		}
	}
	return w.WriteBuf(data)
}

// WriteBuf implements the ResponseWriter.WriteBuf method.
func (w *response) WriteBuf(m []byte) (err error) {
	if m == nil {
		return &Error{Err: "nil message"}
	}
	switch {
	case w.conn._UDP != nil:
		_, err := w.conn._UDP.WriteTo(m, w.conn.remoteAddr)
		if err != nil {
			return err
		}
	case w.conn._TCP != nil:
		if len(m) > MaxMsgSize {
			return ErrBuf
		}
		l := make([]byte, 2)
		l[0], l[1] = packUint16(uint16(len(m)))
		n, err := w.conn._TCP.Write(l)
		if err != nil {
			return err
		}
		if n != 2 {
			return io.ErrShortWrite
		}
		n, err = w.conn._TCP.Write(m)
		if err != nil {
			return err
		}
		i := n
		if i < len(m) {
			j, err := w.conn._TCP.Write(m[i:len(m)])
			if err != nil {
				return err
			}
			i += j
		}
		n = i
	}
	return nil
}

// RemoteAddr implements the ResponseWriter.RemoteAddr method.
func (w *response) RemoteAddr() net.Addr { return w.conn.remoteAddr }

// TsigStatus implements the ResponseWriter.TsigStatus method.
func (w *response) TsigStatus() error { return w.tsigStatus }

// TsigTimersOnly implements the ResponseWriter.TsigTimersOnly method.
func (w *response) TsigTimersOnly(b bool) { w.tsigTimersOnly = b }
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"

	"github.com/flynn/go-discover/discover"
	lornec "github.com/flynn/lorne/client"
	"github.com/flynn/lorne/types"
	sampic "github.com/flynn/sampi/client"
	"github.com/flynn/sampi/types"
	"github.com/titanous/go-dockerclient"
	"github.com/titanous/go-tigertonic"
)

func main() {
	var err error
	scheduler, err = sampic.New()
	if err != nil {
		log.Fatal(err)
	}
	disc, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	mux := tigertonic.NewTrieServeMux()
	mux.Handle("POST", "/apps/{app_id}/formation/{formation_id}", tigertonic.Marshaled(changeFormation))
	mux.Handle("GET", "/apps/{app_id}/jobs", tigertonic.Marshaled(getJobs))
	mux.HandleFunc("GET", "/apps/{app_id}/jobs/{job_id}/logs", getJobLog)
	mux.HandleFunc("POST", "/apps/{app_id}/jobs", runJob)
	http.ListenAndServe("127.0.0.1:1200", tigertonic.Logged(mux, nil))
}

var scheduler *sampic.Client
var disc *discover.Client

type Job struct {
	ID   string `json:"id"`
	Type string `json:"type"`
}

// GET /apps/{app_id}/jobs
func getJobs(u *url.URL, h http.Header) (int, http.Header, []Job, error) {
	state, err := scheduler.State()
	if err != nil {
		return 500, nil, nil, err
	}

	q := u.Query()
	prefix := q.Get("app_id") + "-"
	jobs := make([]Job, 0)
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				typ := strings.Split(job.ID[len(prefix):], ".")[0]
				jobs = append(jobs, Job{ID: job.ID, Type: typ})
			}
		}
	}

	return 200, nil, jobs, nil
}

type Formation struct {
	Quantity int    `json:"quantity"`
	Type     string `json:"type"`
}

func shelfURL() string {
	set, _ := disc.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	return addrs[0]
}

// POST /apps/{app_id}/formation/{formation_id}
func changeFormation(u *url.URL, h http.Header, req *Formation) (int, http.Header, *Formation, error) {
	state, err := scheduler.State()
	if err != nil {
		log.Println("scheduler state error", err)
		return 500, nil, nil, err
	}

	q := u.Query()
	req.Type = q.Get("formation_id")
	prefix := q.Get("app_id") + "-" + req.Type + "."
	var jobs []*sampi.Job
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				if job.Attributes == nil {
					job.Attributes = make(map[string]string)
				}
				job.Attributes["host_id"] = host.ID
				jobs = append(jobs, job)
			}
		}
	}

	if req.Quantity < 0 {
		req.Quantity = 0
	}
	diff := req.Quantity - len(jobs)
	log.Printf("have %d %s, diff %d", len(jobs), req.Type, diff)
	if diff > 0 {
		config := &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          []string{"start", req.Type},
			AttachStdout: true,
			AttachStderr: true,
			Env:          []string{"SLUG_URL=http://" + shelfURL() + "/" + q.Get("app_id") + ".tgz"},
		}
		schedReq := &sampi.ScheduleReq{
			HostJobs: make(map[string][]*sampi.Job),
		}
	outer:
		for {
			for host := range state {
				schedReq.HostJobs[host] = append(schedReq.HostJobs[host], &sampi.Job{ID: prefix + randomID(), Config: config})
				diff--
				if diff == 0 {
					break outer
				}
			}
		}

		res, err := scheduler.Schedule(schedReq)
		if err != nil || !res.Success {
			log.Println("schedule error", err)
			return 500, nil, nil, err
		}
	} else if diff < 0 {
		for _, job := range jobs[:-diff] {
			host, err := lornec.New(job.Attributes["host_id"])
			if err != nil {
				log.Println("error connecting to", job.Attributes["host_id"], err)
				continue
			}
			if err := host.StopJob(job.ID); err != nil {
				log.Println("error stopping", job.ID, "on", job.Attributes["host_id"], err)
			}
		}
	}

	return 200, nil, req, nil
}

// GET /apps/{app_id}/jobs/{job_id}/logs
func getJobLog(w http.ResponseWriter, req *http.Request) {
	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	q := req.URL.Query()
	jobID := q.Get("job_id")
	if prefix := q.Get("app_id") + "-"; !strings.HasPrefix(jobID, prefix) {
		jobID = prefix + jobID
	}
	var job *sampi.Job
	var host sampi.Host
outer:
	for _, host = range state {
		for _, job = range host.Jobs {
			if job.ID == jobID {
				break outer
			}
		}
		job = nil
	}
	if job == nil {
		w.WriteHeader(404)
		return
	}

	attachReq := &lorne.AttachReq{
		JobID: job.ID,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagLogs,
	}
	err, errChan := lorneAttach(host.ID, attachReq, w, nil)
	if err != nil {
		w.WriteHeader(500)
		log.Println("attach error", err)
		return
	}
	if err := <-errChan; err != nil {
		log.Println("attach failed", err)
	}
}

type NewJob struct {
	Cmd     []string          `json:"cmd"`
	Env     map[string]string `json:"env"`
	Attach  bool              `json:"attach"`
	TTY     bool              `json:"tty"`
	Columns int               `json:"tty_columns"`
	Lines   int               `json:"tty_lines"`
}

// POST /apps/{app_id}/jobs
func runJob(w http.ResponseWriter, req *http.Request) {
	var jobReq NewJob
	if err := json.NewDecoder(req.Body).Decode(&jobReq); err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}
	// pick a random host
	var hostID string
	for hostID = range state {
		break
	}
	if hostID == "" {
		w.WriteHeader(500)
		log.Println("no hosts found")
		return
	}

	env := make([]string, 0, len(jobReq.Env))
	for k, v := range jobReq.Env {
		env = append(env, k+"="+v)
	}

	q := req.URL.Query()
	job := &sampi.Job{
		ID: q.Get("app_id") + "-run." + randomID(),
		Config: &docker.Config{
			Image:        "ubuntu",
			Cmd:          jobReq.Cmd,
			AttachStdin:  true,
			AttachStdout: true,
			AttachStderr: true,
			StdinOnce:    true,
			Env:          env,
		},
	}
	if jobReq.TTY {
		job.Config.Tty = true
	}
	if jobReq.Attach {
		job.Config.AttachStdin = true
		job.Config.StdinOnce = true
		job.Config.OpenStdin = true
	}

	outR, outW := io.Pipe()
	inR, inW := io.Pipe()
	defer outR.Close()
	defer inW.Close()
	var errChan <-chan error
	if jobReq.Attach {
		attachReq := &lorne.AttachReq{
			JobID:  job.ID,
			Flags:  lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
			Height: jobReq.Lines,
			Width:  jobReq.Columns,
		}
		err, errChan = lorneAttach(hostID, attachReq, outW, inR)
		if err != nil {
			w.WriteHeader(500)
			log.Println("attach failed", err)
			return
		}
	}

	res, err := scheduler.Schedule(&sampi.ScheduleReq{HostJobs: map[string][]*sampi.Job{hostID: {job}}})
	if err != nil || !res.Success {
		w.WriteHeader(500)
		log.Println("schedule failed", err)
		return
	}

	if jobReq.Attach {
		w.Header().Set("Content-Type", "application/vnd.flynn.hijack")
		w.Header().Set("Content-Length", "0")
		w.WriteHeader(200)
		conn, bufrw, err := w.(http.Hijacker).Hijack()
		if err != nil {
			return
		}
		bufrw.Flush()
		go func() {
			buf := make([]byte, bufrw.Reader.Buffered())
			bufrw.Read(buf)
			inW.Write(buf)
			io.Copy(inW, conn)
			inW.Close()
		}()
		go io.Copy(conn, outR)
		<-errChan
		conn.Close()
		return
	}
	w.WriteHeader(200)
}

func lorneAttach(host string, req *lorne.AttachReq, out io.Writer, in io.Reader) (error, <-chan error) {
	services, err := disc.Services("flynn-lorne-attach." + host)
	if err != nil {
		return err, nil
	}
	addrs := services.OnlineAddrs()
	if len(addrs) == 0 {
		return err, nil
	}
	conn, err := net.Dial("tcp", addrs[0])
	if err != nil {
		return err, nil
	}
	err = gob.NewEncoder(conn).Encode(req)
	if err != nil {
		conn.Close()
		return err, nil
	}

	errChan := make(chan error)

	attach := func() {
		defer conn.Close()
		inErr := make(chan error, 1)
		if in != nil {
			go func() {
				io.Copy(conn, in)
			}()
		} else {
			close(inErr)
		}
		_, outErr := io.Copy(out, conn)
		if outErr != nil {
			errChan <- outErr
			return
		}
		errChan <- <-inErr
	}

	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		conn.Close()
		return err, nil
	}
	switch attachState[0] {
	case lorne.AttachError:
		errBytes, err := ioutil.ReadAll(conn)
		conn.Close()
		if err != nil {
			return err, nil
		}
		return errors.New(string(errBytes)), nil
	case lorne.AttachWaiting:
		go func() {
			if _, err := conn.Read(attachState); err != nil {
				conn.Close()
				errChan <- err
				return
			}
			if attachState[0] == lorne.AttachError {
				errBytes, err := ioutil.ReadAll(conn)
				conn.Close()
				if err != nil {
					errChan <- err
					return
				}
				errChan <- errors.New(string(errBytes))
				return
			}
			attach()
		}()
		return nil, errChan
	default:
		go attach()
		return nil, errChan
	}
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<commit_msg>Request TCP port<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"

	"github.com/flynn/go-discover/discover"
	lornec "github.com/flynn/lorne/client"
	"github.com/flynn/lorne/types"
	sampic "github.com/flynn/sampi/client"
	"github.com/flynn/sampi/types"
	"github.com/titanous/go-dockerclient"
	"github.com/titanous/go-tigertonic"
)

func main() {
	var err error
	scheduler, err = sampic.New()
	if err != nil {
		log.Fatal(err)
	}
	disc, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	mux := tigertonic.NewTrieServeMux()
	mux.Handle("POST", "/apps/{app_id}/formation/{formation_id}", tigertonic.Marshaled(changeFormation))
	mux.Handle("GET", "/apps/{app_id}/jobs", tigertonic.Marshaled(getJobs))
	mux.HandleFunc("GET", "/apps/{app_id}/jobs/{job_id}/logs", getJobLog)
	mux.HandleFunc("POST", "/apps/{app_id}/jobs", runJob)
	http.ListenAndServe("127.0.0.1:1200", tigertonic.Logged(mux, nil))
}

var scheduler *sampic.Client
var disc *discover.Client

type Job struct {
	ID   string `json:"id"`
	Type string `json:"type"`
}

// GET /apps/{app_id}/jobs
func getJobs(u *url.URL, h http.Header) (int, http.Header, []Job, error) {
	state, err := scheduler.State()
	if err != nil {
		return 500, nil, nil, err
	}

	q := u.Query()
	prefix := q.Get("app_id") + "-"
	jobs := make([]Job, 0)
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				typ := strings.Split(job.ID[len(prefix):], ".")[0]
				jobs = append(jobs, Job{ID: job.ID, Type: typ})
			}
		}
	}

	return 200, nil, jobs, nil
}

type Formation struct {
	Quantity int    `json:"quantity"`
	Type     string `json:"type"`
}

func shelfURL() string {
	set, _ := disc.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	return addrs[0]
}

// POST /apps/{app_id}/formation/{formation_id}
func changeFormation(u *url.URL, h http.Header, req *Formation) (int, http.Header, *Formation, error) {
	state, err := scheduler.State()
	if err != nil {
		log.Println("scheduler state error", err)
		return 500, nil, nil, err
	}

	q := u.Query()
	req.Type = q.Get("formation_id")
	prefix := q.Get("app_id") + "-" + req.Type + "."
	var jobs []*sampi.Job
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				if job.Attributes == nil {
					job.Attributes = make(map[string]string)
				}
				job.Attributes["host_id"] = host.ID
				jobs = append(jobs, job)
			}
		}
	}

	if req.Quantity < 0 {
		req.Quantity = 0
	}
	diff := req.Quantity - len(jobs)
	log.Printf("have %d %s, diff %d", len(jobs), req.Type, diff)
	if diff > 0 {
		config := &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          []string{"start", req.Type},
			AttachStdout: true,
			AttachStderr: true,
			Env:          []string{"SLUG_URL=http://" + shelfURL() + "/" + q.Get("app_id") + ".tgz"},
		}
		schedReq := &sampi.ScheduleReq{
			HostJobs: make(map[string][]*sampi.Job),
		}
	outer:
		for {
			for host := range state {
				schedReq.HostJobs[host] = append(schedReq.HostJobs[host], &sampi.Job{ID: prefix + randomID(), TCPPorts: 1, Config: config})
				diff--
				if diff == 0 {
					break outer
				}
			}
		}

		res, err := scheduler.Schedule(schedReq)
		if err != nil || !res.Success {
			log.Println("schedule error", err)
			return 500, nil, nil, err
		}
	} else if diff < 0 {
		for _, job := range jobs[:-diff] {
			host, err := lornec.New(job.Attributes["host_id"])
			if err != nil {
				log.Println("error connecting to", job.Attributes["host_id"], err)
				continue
			}
			if err := host.StopJob(job.ID); err != nil {
				log.Println("error stopping", job.ID, "on", job.Attributes["host_id"], err)
			}
		}
	}

	return 200, nil, req, nil
}

// GET /apps/{app_id}/jobs/{job_id}/logs
func getJobLog(w http.ResponseWriter, req *http.Request) {
	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	q := req.URL.Query()
	jobID := q.Get("job_id")
	if prefix := q.Get("app_id") + "-"; !strings.HasPrefix(jobID, prefix) {
		jobID = prefix + jobID
	}
	var job *sampi.Job
	var host sampi.Host
outer:
	for _, host = range state {
		for _, job = range host.Jobs {
			if job.ID == jobID {
				break outer
			}
		}
		job = nil
	}
	if job == nil {
		w.WriteHeader(404)
		return
	}

	attachReq := &lorne.AttachReq{
		JobID: job.ID,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagLogs,
	}
	err, errChan := lorneAttach(host.ID, attachReq, w, nil)
	if err != nil {
		w.WriteHeader(500)
		log.Println("attach error", err)
		return
	}
	if err := <-errChan; err != nil {
		log.Println("attach failed", err)
	}
}

type NewJob struct {
	Cmd     []string          `json:"cmd"`
	Env     map[string]string `json:"env"`
	Attach  bool              `json:"attach"`
	TTY     bool              `json:"tty"`
	Columns int               `json:"tty_columns"`
	Lines   int               `json:"tty_lines"`
}

// POST /apps/{app_id}/jobs
func runJob(w http.ResponseWriter, req *http.Request) {
	var jobReq NewJob
	if err := json.NewDecoder(req.Body).Decode(&jobReq); err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}
	// pick a random host
	var hostID string
	for hostID = range state {
		break
	}
	if hostID == "" {
		w.WriteHeader(500)
		log.Println("no hosts found")
		return
	}

	env := make([]string, 0, len(jobReq.Env))
	for k, v := range jobReq.Env {
		env = append(env, k+"="+v)
	}

	q := req.URL.Query()
	job := &sampi.Job{
		ID: q.Get("app_id") + "-run." + randomID(),
		Config: &docker.Config{
			Image:        "ubuntu",
			Cmd:          jobReq.Cmd,
			AttachStdin:  true,
			AttachStdout: true,
			AttachStderr: true,
			StdinOnce:    true,
			Env:          env,
		},
	}
	if jobReq.TTY {
		job.Config.Tty = true
	}
	if jobReq.Attach {
		job.Config.AttachStdin = true
		job.Config.StdinOnce = true
		job.Config.OpenStdin = true
	}

	outR, outW := io.Pipe()
	inR, inW := io.Pipe()
	defer outR.Close()
	defer inW.Close()
	var errChan <-chan error
	if jobReq.Attach {
		attachReq := &lorne.AttachReq{
			JobID:  job.ID,
			Flags:  lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
			Height: jobReq.Lines,
			Width:  jobReq.Columns,
		}
		err, errChan = lorneAttach(hostID, attachReq, outW, inR)
		if err != nil {
			w.WriteHeader(500)
			log.Println("attach failed", err)
			return
		}
	}

	res, err := scheduler.Schedule(&sampi.ScheduleReq{HostJobs: map[string][]*sampi.Job{hostID: {job}}})
	if err != nil || !res.Success {
		w.WriteHeader(500)
		log.Println("schedule failed", err)
		return
	}

	if jobReq.Attach {
		w.Header().Set("Content-Type", "application/vnd.flynn.hijack")
		w.Header().Set("Content-Length", "0")
		w.WriteHeader(200)
		conn, bufrw, err := w.(http.Hijacker).Hijack()
		if err != nil {
			return
		}
		bufrw.Flush()
		go func() {
			buf := make([]byte, bufrw.Reader.Buffered())
			bufrw.Read(buf)
			inW.Write(buf)
			io.Copy(inW, conn)
			inW.Close()
		}()
		go io.Copy(conn, outR)
		<-errChan
		conn.Close()
		return
	}
	w.WriteHeader(200)
}

func lorneAttach(host string, req *lorne.AttachReq, out io.Writer, in io.Reader) (error, <-chan error) {
	services, err := disc.Services("flynn-lorne-attach." + host)
	if err != nil {
		return err, nil
	}
	addrs := services.OnlineAddrs()
	if len(addrs) == 0 {
		return err, nil
	}
	conn, err := net.Dial("tcp", addrs[0])
	if err != nil {
		return err, nil
	}
	err = gob.NewEncoder(conn).Encode(req)
	if err != nil {
		conn.Close()
		return err, nil
	}

	errChan := make(chan error)

	attach := func() {
		defer conn.Close()
		inErr := make(chan error, 1)
		if in != nil {
			go func() {
				io.Copy(conn, in)
			}()
		} else {
			close(inErr)
		}
		_, outErr := io.Copy(out, conn)
		if outErr != nil {
			errChan <- outErr
			return
		}
		errChan <- <-inErr
	}

	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		conn.Close()
		return err, nil
	}
	switch attachState[0] {
	case lorne.AttachError:
		errBytes, err := ioutil.ReadAll(conn)
		conn.Close()
		if err != nil {
			return err, nil
		}
		return errors.New(string(errBytes)), nil
	case lorne.AttachWaiting:
		go func() {
			if _, err := conn.Read(attachState); err != nil {
				conn.Close()
				errChan <- err
				return
			}
			if attachState[0] == lorne.AttachError {
				errBytes, err := ioutil.ReadAll(conn)
				conn.Close()
				if err != nil {
					errChan <- err
					return
				}
				errChan <- errors.New(string(errBytes))
				return
			}
			attach()
		}()
		return nil, errChan
	default:
		go attach()
		return nil, errChan
	}
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<|endoftext|>


<commit_before>package main

import (
  "flag"

  "github.com/r4start/web-tracer/tracer"
)

func getServeAddress() tracer.ServerParameters {
  var params tracer.ServerParameters

  flag.StringVar(&params.Host, "host", "localhost", "IP address for listening")
  flag.StringVar(&params.Port, "port", "80", "Port number")
  flag.StringVar(&params.DbName, "dbname", "tracer.db", "Database name or path")
  flag.StringVar(&params.SiteRoot, "site-root", "", "Path to site root folder")

  flag.Parse()

  return params
}

func main() {
  params := getServeAddress()

  tracer.StartServer(params)
}
<commit_msg>By default server will use in-memory database.<commit_after>package main

import (
  "flag"

  "github.com/r4start/web-tracer/tracer"
)

func getServeAddress() tracer.ServerParameters {
  var params tracer.ServerParameters

  flag.StringVar(&params.Host, "host", "localhost", "IP address for listening")
  flag.StringVar(&params.Port, "port", "80", "Port number")
  flag.StringVar(&params.DbName, "dbname", ":memory:", "Database name or path")
  flag.StringVar(&params.SiteRoot, "site-root", "", "Path to site root folder")

  flag.Parse()

  return params
}

func main() {
  params := getServeAddress()

  tracer.StartServer(params)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

const MAX_NAME_LENGTH = 32
const HISTORY_LEN = 20

var RE_STRIP_NAME = regexp.MustCompile("[^0-9A-Za-z_]")

type Clients map[string]*Client

type Server struct {
	sshConfig *ssh.ServerConfig
	done      chan struct{}
	clients   Clients
	lock      sync.Mutex
	count     int
	history   *History
	admins    map[string]struct{}   // fingerprint lookup
	banned    map[string]*time.Time // fingerprint lookup
}

func NewServer(privateKey []byte) (*Server, error) {
	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	server := Server{
		done:    make(chan struct{}),
		clients: Clients{},
		count:   0,
		history: NewHistory(HISTORY_LEN),
		admins:  map[string]struct{}{},
		banned:  map[string]*time.Time{},
	}

	config := ssh.ServerConfig{
		NoClientAuth: false,
		// Auth-related things should be constant-time to avoid timing attacks.
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			fingerprint := Fingerprint(key)
			if server.IsBanned(fingerprint) {
				return nil, fmt.Errorf("Banned.")
			}
			perm := &ssh.Permissions{Extensions: map[string]string{"fingerprint": fingerprint}}
			return perm, nil
		},
	}
	config.AddHostKey(signer)

	server.sshConfig = &config

	return &server, nil
}

func (s *Server) Len() int {
	return len(s.clients)
}

func (s *Server) Broadcast(msg string, except *Client) {
	logger.Debugf("Broadcast to %d: %s", s.Len(), msg)
	s.history.Add(msg)

	for _, client := range s.clients {
		if except != nil && client == except {
			continue
		}
		client.Msg <- msg
	}
}

func (s *Server) Add(client *Client) {
	go func() {
		client.WriteLines(s.history.Get(10))
		client.Write(fmt.Sprintf("-> Welcome to ssh-chat. Enter /help for more."))
	}()

	s.lock.Lock()
	s.count++

	newName, err := s.proposeName(client.Name)
	if err != nil {
		client.Msg <- fmt.Sprintf("-> Your name '%s' is not available, renamed to '%s'. Use /nick <name> to change it.", client.Name, newName)
	}

	client.Rename(newName)
	s.clients[client.Name] = client
	num := len(s.clients)
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)", client.Name, num), client)
}

func (s *Server) Remove(client *Client) {
	s.lock.Lock()
	delete(s.clients, client.Name)
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s left.", client.Name), nil)
}

func (s *Server) proposeName(name string) (string, error) {
	// Assumes caller holds lock.
	var err error
	name = RE_STRIP_NAME.ReplaceAllString(name, "")

	if len(name) > MAX_NAME_LENGTH {
		name = name[:MAX_NAME_LENGTH]
	} else if len(name) == 0 {
		name = fmt.Sprintf("Guest%d", s.count)
	}

	_, collision := s.clients[name]
	if collision {
		err = fmt.Errorf("%s is not available.", name)
		name = fmt.Sprintf("Guest%d", s.count)
	}

	return name, err
}

func (s *Server) Rename(client *Client, newName string) {
	s.lock.Lock()

	newName, err := s.proposeName(newName)
	if err != nil {
		client.Msg <- fmt.Sprintf("-> %s", err)
		s.lock.Unlock()
		return
	}

	// TODO: Use a channel/goroutine for adding clients, rathern than locks?
	delete(s.clients, client.Name)
	oldName := client.Name
	client.Rename(newName)
	s.clients[client.Name] = client
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s is now known as %s.", oldName, newName), nil)
}

func (s *Server) List(prefix *string) []string {
	r := []string{}

	for name, _ := range s.clients {
		if prefix != nil && !strings.HasPrefix(name, *prefix) {
			continue
		}
		r = append(r, name)
	}

	return r
}

func (s *Server) Who(name string) *Client {
	return s.clients[name]
}

func (s *Server) Op(fingerprint string) {
	logger.Infof("Adding admin: %s", fingerprint)
	s.lock.Lock()
	s.admins[fingerprint] = struct{}{}
	s.lock.Unlock()
}

func (s *Server) IsOp(client *Client) bool {
	_, r := s.admins[client.Fingerprint()]
	return r
}

func (s *Server) IsBanned(fingerprint string) bool {
	ban, hasBan := s.banned[fingerprint]
	if !hasBan {
		return false
	}
	if ban == nil {
		return true
	}
	if ban.Before(time.Now()) {
		s.Unban(fingerprint)
		return false
	}
	return true
}

func (s *Server) Ban(fingerprint string, duration *time.Duration) {
	var until *time.Time
	s.lock.Lock()
	if duration != nil {
		when := time.Now().Add(*duration)
		until = &when
	}
	s.banned[fingerprint] = until
	s.lock.Unlock()
}

func (s *Server) Unban(fingerprint string) {
	s.lock.Lock()
	delete(s.banned, fingerprint)
	s.lock.Unlock()
}

func (s *Server) Start(laddr string) error {
	// Once a ServerConfig has been configured, connections can be
	// accepted.
	socket, err := net.Listen("tcp", laddr)
	if err != nil {
		return err
	}

	logger.Infof("Listening on %s", laddr)

	go func() {
		for {
			conn, err := socket.Accept()

			if err != nil {
				// TODO: Handle shutdown more gracefully?
				logger.Errorf("Failed to accept connection, aborting loop: %v", err)
				return
			}

			// Goroutineify to resume accepting sockets early.
			go func() {
				// From a standard TCP connection to an encrypted SSH connection
				sshConn, channels, requests, err := ssh.NewServerConn(conn, s.sshConfig)
				if err != nil {
					logger.Errorf("Failed to handshake: %v", err)
					return
				}

				logger.Infof("Connection #%d from: %s, %s, %s", s.count+1, sshConn.RemoteAddr(), sshConn.User(), sshConn.ClientVersion())

				go ssh.DiscardRequests(requests)

				client := NewClient(s, sshConn)
				go client.handleChannels(channels)
			}()
		}
	}()

	go func() {
		<-s.done
		socket.Close()
	}()

	return nil
}

func (s *Server) Stop() {
	for _, client := range s.clients {
		client.Conn.Close()
	}

	close(s.done)
}

func Fingerprint(k ssh.PublicKey) string {
	hash := md5.Sum(k.Marshal())
	r := fmt.Sprintf("% x", hash)
	return strings.Replace(r, " ", ":", -1)
}
<commit_msg>More evil inputs.<commit_after>package main

import (
	"crypto/md5"
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"

	"code.google.com/p/rog-go/new9p/client"

	"golang.org/x/crypto/ssh"
)

const MAX_NAME_LENGTH = 32
const HISTORY_LEN = 20

var RE_STRIP_NAME = regexp.MustCompile("[^0-9A-Za-z_]")

type Clients map[string]*Client

type Server struct {
	sshConfig *ssh.ServerConfig
	done      chan struct{}
	clients   Clients
	lock      sync.Mutex
	count     int
	history   *History
	admins    map[string]struct{}   // fingerprint lookup
	banned    map[string]*time.Time // fingerprint lookup
}

func NewServer(privateKey []byte) (*Server, error) {
	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	server := Server{
		done:    make(chan struct{}),
		clients: Clients{},
		count:   0,
		history: NewHistory(HISTORY_LEN),
		admins:  map[string]struct{}{},
		banned:  map[string]*time.Time{},
	}

	config := ssh.ServerConfig{
		NoClientAuth: false,
		// Auth-related things should be constant-time to avoid timing attacks.
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			fingerprint := Fingerprint(key)
			if server.IsBanned(fingerprint) {
				return nil, fmt.Errorf("Banned.")
			}
			perm := &ssh.Permissions{Extensions: map[string]string{"fingerprint": fingerprint}}
			return perm, nil
		},
	}
	config.AddHostKey(signer)

	server.sshConfig = &config

	return &server, nil
}

func (s *Server) Len() int {
	return len(s.clients)
}

func (s *Server) Broadcast(msg string, except *Client) {
	logger.Debugf("Broadcast to %d: %s", s.Len(), msg)
	s.history.Add(msg)

	for _, client := range s.clients {
		if except != nil && client == except {
			continue
		}
		client.Msg <- msg
	}
}

func (s *Server) Add(client *Client) {
	go func() {
		client.WriteLines(s.history.Get(10))
		client.Write(fmt.Sprintf("-> Welcome to ssh-chat. Enter /help for more."))
	}()

	s.lock.Lock()
	s.count++

	newName, err := s.proposeName(client.Name)
	if err != nil {
		client.Msg <- fmt.Sprintf("-> Your name '%s' is not available, renamed to '%s'. Use /nick <name> to change it.", client.Name, newName)
	}

	client.Rename(newName)
	s.clients[client.Name] = client
	num := len(s.clients)
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)", client.Name, num), client)
}

func (s *Server) Remove(client *Client) {
	s.lock.Lock()
	delete(s.clients, client.Name)
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s left.", client.Name), nil)
}

func (s *Server) proposeName(name string) (string, error) {
	// Assumes caller holds lock.
	var err error
	name = RE_STRIP_NAME.ReplaceAllString(name, "")

	if len(name) > MAX_NAME_LENGTH {
		name = name[:MAX_NAME_LENGTH]
	} else if len(name) == 0 {
		name = fmt.Sprintf("Guest%d", s.count)
	}

	_, collision := s.clients[name]
	if collision {
		err = fmt.Errorf("%s is not available.", name)
		name = fmt.Sprintf("Guest%d", s.count)
	}

	return name, err
}

func (s *Server) Rename(client *Client, newName string) {
	s.lock.Lock()

	newName, err := s.proposeName(newName)
	if err != nil {
		client.Msg <- fmt.Sprintf("-> %s", err)
		s.lock.Unlock()
		return
	}

	// TODO: Use a channel/goroutine for adding clients, rathern than locks?
	delete(s.clients, client.Name)
	oldName := client.Name
	client.Rename(newName)
	s.clients[client.Name] = client
	s.lock.Unlock()

	s.Broadcast(fmt.Sprintf("* %s is now known as %s.", oldName, newName), nil)
}

func (s *Server) List(prefix *string) []string {
	r := []string{}

	for name, _ := range s.clients {
		if prefix != nil && !strings.HasPrefix(name, *prefix) {
			continue
		}
		r = append(r, name)
	}

	return r
}

func (s *Server) Who(name string) *Client {
	return s.clients[name]
}

func (s *Server) Op(fingerprint string) {
	logger.Infof("Adding admin: %s", fingerprint)
	s.lock.Lock()
	s.admins[fingerprint] = struct{}{}
	s.lock.Unlock()
}

func (s *Server) IsOp(client *Client) bool {
	_, r := s.admins[client.Fingerprint()]
	return r
}

func (s *Server) IsBanned(fingerprint string) bool {
	ban, hasBan := s.banned[fingerprint]
	if !hasBan {
		return false
	}
	if ban == nil {
		return true
	}
	if ban.Before(time.Now()) {
		s.Unban(fingerprint)
		return false
	}
	return true
}

func (s *Server) Ban(fingerprint string, duration *time.Duration) {
	var until *time.Time
	s.lock.Lock()
	if duration != nil {
		when := time.Now().Add(*duration)
		until = &when
	}
	s.banned[fingerprint] = until
	s.lock.Unlock()
}

func (s *Server) Unban(fingerprint string) {
	s.lock.Lock()
	delete(s.banned, fingerprint)
	s.lock.Unlock()
}

func (s *Server) Start(laddr string) error {
	// Once a ServerConfig has been configured, connections can be
	// accepted.
	socket, err := net.Listen("tcp", laddr)
	if err != nil {
		return err
	}

	logger.Infof("Listening on %s", laddr)

	go func() {
		for {
			conn, err := socket.Accept()

			if err != nil {
				// TODO: Handle shutdown more gracefully?
				logger.Errorf("Failed to accept connection, aborting loop: %v", err)
				return
			}

			// Goroutineify to resume accepting sockets early.
			go func() {
				// From a standard TCP connection to an encrypted SSH connection
				sshConn, channels, requests, err := ssh.NewServerConn(conn, s.sshConfig)
				if err != nil {
					logger.Errorf("Failed to handshake: %v", err)
					return
				}

				version := client.Conn.ClientVersion()
				if len(version) > 100 {
					version = []byte("Evil Jerk with a superlong string")
				}
				logger.Infof("Connection #%d from: %s, %s, %s", s.count+1, sshConn.RemoteAddr(), sshConn.User(), version)

				go ssh.DiscardRequests(requests)

				client := NewClient(s, sshConn)
				go client.handleChannels(channels)
			}()
		}
	}()

	go func() {
		<-s.done
		socket.Close()
	}()

	return nil
}

func (s *Server) Stop() {
	for _, client := range s.clients {
		client.Conn.Close()
	}

	close(s.done)
}

func Fingerprint(k ssh.PublicKey) string {
	hash := md5.Sum(k.Marshal())
	r := fmt.Sprintf("% x", hash)
	return strings.Replace(r, " ", ":", -1)
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(*http.Request) CodecRequest
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads request and returns the RPC method name.
	Method() (string, error)
	// Reads request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes response using the RPC method reply. The error parameter is
	// the error returned by the method call, if any.
	WriteResponse(http.ResponseWriter, interface{}, error) error
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer() *Server {
	return &Server{
		codecs:   make(map[string]Codec),
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	codecs   map[string]Codec
	services *serviceMap
}

// RegisterCodec adds a new codec to the server.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (s *Server) RegisterCodec(codec Codec, contentType string) {
	s.codecs[strings.ToLower(contentType)] = codec
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has three arguments: *http.Request, *args, *reply.
//    - All three arguments are pointers.
//    - The second and third arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	return s.services.register(receiver, name)
}

// HasMethod returns true if the given method is registered.
//
// The method uses a dotted notation as in "Service.Method".
func (s *Server) HasMethod(method string) bool {
	if _, _, err := s.services.get(method); err == nil {
		return true
	}
	return false
}

// ServeHTTP
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, 405, "rpc: POST method required, received "+r.Method)
		return
	}
	contentType := r.Header.Get("Content-Type")
	idx := strings.Index(contentType, ";")
	if idx != -1 {
		contentType = contentType[:idx]
	}
	codec := s.codecs[strings.ToLower(contentType)]
	if codec == nil {
		writeError(w, 415, "rpc: unrecognized Content-Type: "+contentType)
		return
	}
	// Create a new codec request.
	codecReq := codec.NewRequest(r)
	// Get service method to be called.
	method, errMethod := codecReq.Method()
	if errMethod != nil {
		writeError(w, 400, errMethod.Error())
		return
	}
	serviceSpec, methodSpec, errGet := s.services.get(method)
	if errGet != nil {
		writeError(w, 400, errGet.Error())
		return
	}
	// Decode the args.
	args := reflect.New(methodSpec.argsType)
	if methodSpec.argsType.NumField() > 0 {
		if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
			writeError(w, 400, errRead.Error())
			return
		}
	}
	// Call the service method.
	reply := reflect.New(methodSpec.replyType)
	errValue := methodSpec.method.Func.Call([]reflect.Value{
		serviceSpec.rcvr,
		reflect.ValueOf(r),
		args,
		reply,
	})
	// Cast the result to error if needed.
	var errResult error
	errInter := errValue[0].Interface()
	if errInter != nil {
		errResult = errInter.(error)
	}
	// Prevents Internet Explorer from MIME-sniffing a response away
	// from the declared content-type
	w.Header().Set("x-content-type-options", "nosniff")
	// Encode the response.
	if errWrite := codecReq.WriteResponse(w, reply.Interface(), errResult); errWrite != nil {
		writeError(w, 400, errWrite.Error())
	}
}

func writeError(w http.ResponseWriter, status int, msg string) {
	w.WriteHeader(status)
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprint(w, msg)
}
<commit_msg>revert server.go changes<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(*http.Request) CodecRequest
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads request and returns the RPC method name.
	Method() (string, error)
	// Reads request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes response using the RPC method reply. The error parameter is
	// the error returned by the method call, if any.
	WriteResponse(http.ResponseWriter, interface{}, error) error
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer() *Server {
	return &Server{
		codecs:   make(map[string]Codec),
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	codecs   map[string]Codec
	services *serviceMap
}

// RegisterCodec adds a new codec to the server.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (s *Server) RegisterCodec(codec Codec, contentType string) {
	s.codecs[strings.ToLower(contentType)] = codec
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has three arguments: *http.Request, *args, *reply.
//    - All three arguments are pointers.
//    - The second and third arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	return s.services.register(receiver, name)
}

// HasMethod returns true if the given method is registered.
//
// The method uses a dotted notation as in "Service.Method".
func (s *Server) HasMethod(method string) bool {
	if _, _, err := s.services.get(method); err == nil {
		return true
	}
	return false
}

// ServeHTTP
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, 405, "rpc: POST method required, received "+r.Method)
		return
	}
	contentType := r.Header.Get("Content-Type")
	idx := strings.Index(contentType, ";")
	if idx != -1 {
		contentType = contentType[:idx]
	}
	codec := s.codecs[strings.ToLower(contentType)]
	if codec == nil {
		writeError(w, 415, "rpc: unrecognized Content-Type: "+contentType)
		return
	}
	// Create a new codec request.
	codecReq := codec.NewRequest(r)
	// Get service method to be called.
	method, errMethod := codecReq.Method()
	if errMethod != nil {
		writeError(w, 400, errMethod.Error())
		return
	}
	serviceSpec, methodSpec, errGet := s.services.get(method)
	if errGet != nil {
		writeError(w, 400, errGet.Error())
		return
	}
	// Decode the args.
	args := reflect.New(methodSpec.argsType)
	if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
		writeError(w, 400, errRead.Error())
		return
	}
	// Call the service method.
	reply := reflect.New(methodSpec.replyType)
	errValue := methodSpec.method.Func.Call([]reflect.Value{
		serviceSpec.rcvr,
		reflect.ValueOf(r),
		args,
		reply,
	})
	// Cast the result to error if needed.
	var errResult error
	errInter := errValue[0].Interface()
	if errInter != nil {
		errResult = errInter.(error)
	}
	// Prevents Internet Explorer from MIME-sniffing a response away
	// from the declared content-type
	w.Header().Set("x-content-type-options", "nosniff")
	// Encode the response.
	if errWrite := codecReq.WriteResponse(w, reply.Interface(), errResult); errWrite != nil {
		writeError(w, 400, errWrite.Error())
	}
}

func writeError(w http.ResponseWriter, status int, msg string) {
	w.WriteHeader(status)
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprint(w, msg)
}
<|endoftext|>


<commit_before>package tftp

import (
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

// NewServer creates TFTP server. It requires two functions to handle
// read and write requests.
// In case nil is provided for read or write handler the respective
// operation is disabled.
func NewServer(readHandler func(filename string, rf io.ReaderFrom) error,
	writeHandler func(filename string, wt io.WriterTo) error) *Server {
	return &Server{
		readHandler:  readHandler,
		writeHandler: writeHandler,
		timeout:      defaultTimeout,
		retries:      defaultRetries,
	}
}

// RequestPacketInfo provides a method of getting the local IP address
// that is handling a UDP request.  It relies for its accuracy on the
// OS providing methods to inspect the underlying UDP and IP packets
// directly.
type RequestPacketInfo interface {
	// LocalAddr returns the IP address we are servicing the request on.
	// If it is unable to determine what address that is, the returned
	// net.IP will be nil.
	LocalIP() net.IP
}

type Server struct {
	readHandler  func(filename string, rf io.ReaderFrom) error
	writeHandler func(filename string, wt io.WriterTo) error
	backoff      backoffFunc
	conn         *net.UDPConn
	quit         chan chan struct{}
	wg           sync.WaitGroup
	timeout      time.Duration
	retries      int
	maxBlockLen  int
	sendAEnable  bool /* senderAnticipate enable by server */
	sendAWinSz   uint
}

// SetAnticipate provides an experimental feature in which when a packets
// is requested the server will keep sending a number of packets before
// checking whether an ack has been received. It improves tftp downloading
// speed by a few times.
// The argument winsz specifies how many packets will be sent before
// waiting for an ack packet.
// When winsz is bigger than 1, the feature is enabled, and the server
// runs through a different experimental code path. When winsz is 0 or 1,
// the feature is disabled.
func (s *Server) SetAnticipate(winsz uint) {
	if winsz > 1 {
		s.sendAEnable = true
		s.sendAWinSz = winsz
	} else {
		s.sendAEnable = false
		s.sendAWinSz = 1
	}
}

// SetTimeout sets maximum time server waits for single network
// round-trip to succeed.
// Default is 5 seconds.
func (s *Server) SetTimeout(t time.Duration) {
	if t <= 0 {
		s.timeout = defaultTimeout
	} else {
		s.timeout = t
	}
}

// SetBlockSize sets the maximum size of an individual data block.
// This must be a value between 512 (the default block size for TFTP)
// and 65456 (the max size a UDP packet payload can be).
//
// This is an advisory value -- it will be clamped to the smaller of
// the block size the client wants and the MTU of the interface being
// communicated over munis overhead.
func (s *Server) SetBlockSize(i int) {
	if i > 512 && i < 65465 {
		s.maxBlockLen = i
	}
}

// SetRetries sets maximum number of attempts server made to transmit a
// packet.
// Default is 5 attempts.
func (s *Server) SetRetries(count int) {
	if count < 1 {
		s.retries = defaultRetries
	} else {
		s.retries = count
	}
}

// SetBackoff sets a user provided function that is called to provide a
// backoff duration prior to retransmitting an unacknowledged packet.
func (s *Server) SetBackoff(h backoffFunc) {
	s.backoff = h
}

// ListenAndServe binds to address provided and start the server.
// ListenAndServe returns when Shutdown is called.
func (s *Server) ListenAndServe(addr string) error {
	a, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", a)
	if err != nil {
		return err
	}
	return s.Serve(conn)
}

// Serve starts server provided already opened UDP connecton. It is
// useful for the case when you want to run server in separate goroutine
// but still want to be able to handle any errors opening connection.
// Serve returns when Shutdown is called or connection is closed.
func (s *Server) Serve(conn *net.UDPConn) error {
	defer conn.Close()
	laddr := conn.LocalAddr()
	host, _, err := net.SplitHostPort(laddr.String())
	if err != nil {
		return err
	}
	s.conn = conn
	// Having seperate control paths for IP4 and IP6 is annoying,
	// but necessary at this point.
	addr := net.ParseIP(host)
	if addr == nil {
		return fmt.Errorf("Failed to determine IP class of listening address")
	}
	var conn4 *ipv4.PacketConn
	var conn6 *ipv6.PacketConn
	if addr.To4() != nil {
		conn4 = ipv4.NewPacketConn(conn)
		if err := conn4.SetControlMessage(ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
			conn4 = nil
		}
	} else {
		conn6 = ipv6.NewPacketConn(conn)
		if err := conn6.SetControlMessage(ipv6.FlagDst|ipv6.FlagInterface, true); err != nil {
			conn6 = nil
		}
	}

	s.quit = make(chan chan struct{})
	for {
		select {
		case q := <-s.quit:
			q <- struct{}{}
			return nil
		default:
			var err error
			if conn4 != nil {
				err = s.processRequest4(conn4)
			} else if conn6 != nil {
				err = s.processRequest6(conn6)
			} else {
				err = s.processRequest()
			}
			if err != nil {
				// TODO: add logging handler
			}
		}
	}
}

// Yes, I don't really like having seperate IPv4 and IPv6 variants,
// bit we are relying on the low-level packet control channel info to
// get a reliable source address, and those have different types and
// the struct itself is not easily interface-ized or embedded.
//
// If control is nil for whatever reason (either things not being
// implemented on a target OS or whatever other reason), localIP
// (and hence LocalIP()) will return a nil IP address.
func (s *Server) processRequest4(conn4 *ipv4.PacketConn) error {
	buf := make([]byte, datagramLength)
	cnt, control, srcAddr, err := conn4.ReadFrom(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	maxSz := blockLength
	var localAddr net.IP
	if control != nil {
		localAddr = control.Dst
		if intf, err := net.InterfaceByIndex(control.IfIndex); err == nil {
			// mtu - ipv4 overhead - udp overhead
			maxSz = intf.MTU - 28
		}
	}
	return s.handlePacket(localAddr, srcAddr.(*net.UDPAddr), buf, cnt, maxSz)
}

func (s *Server) processRequest6(conn6 *ipv6.PacketConn) error {
	buf := make([]byte, datagramLength)
	cnt, control, srcAddr, err := conn6.ReadFrom(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	maxSz := blockLength
	var localAddr net.IP
	if control != nil {
		localAddr = control.Dst
		if intf, err := net.InterfaceByIndex(control.IfIndex); err == nil {
			// mtu - ipv6 overhead - udp overhead
			maxSz = intf.MTU - 48
		}
	}
	return s.handlePacket(localAddr, srcAddr.(*net.UDPAddr), buf, cnt, maxSz)
}

// Fallback if we had problems opening a ipv4/6 control channel
func (s *Server) processRequest() error {
	buf := make([]byte, datagramLength)
	cnt, srcAddr, err := s.conn.ReadFromUDP(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	return s.handlePacket(nil, srcAddr, buf, cnt, blockLength)
}

// Shutdown make server stop listening for new requests, allows
// server to finish outstanding transfers and stops server.
func (s *Server) Shutdown() {
	s.conn.Close()
	q := make(chan struct{})
	s.quit <- q
	<-q
	s.wg.Wait()
}

func (s *Server) handlePacket(localAddr net.IP, remoteAddr *net.UDPAddr, buffer []byte, n, maxBlockLen int) error {
	if s.maxBlockLen > 0 && s.maxBlockLen < maxBlockLen {
		maxBlockLen = s.maxBlockLen
	}
	if maxBlockLen < blockLength {
		maxBlockLen = blockLength
	}
	p, err := parsePacket(buffer[:n])
	if err != nil {
		return err
	}
	switch p := p.(type) {
	case pWRQ:
		filename, mode, opts, err := unpackRQ(p)
		if err != nil {
			return fmt.Errorf("unpack WRQ: %v", err)
		}
		//fmt.Printf("got WRQ (filename=%s, mode=%s, opts=%v)\n", filename, mode, opts)
		conn, err := net.ListenUDP("udp", &net.UDPAddr{})
		if err != nil {
			return err
		}
		if err != nil {
			return fmt.Errorf("open transmission: %v", err)
		}
		wt := &receiver{
			send:        make([]byte, datagramLength),
			receive:     make([]byte, datagramLength),
			conn:        conn,
			retry:       &backoff{handler: s.backoff},
			timeout:     s.timeout,
			retries:     s.retries,
			addr:        remoteAddr,
			localIP:     localAddr,
			mode:        mode,
			opts:        opts,
			maxBlockLen: maxBlockLen,
		}
		s.wg.Add(1)
		go func() {
			if s.writeHandler != nil {
				err := s.writeHandler(filename, wt)
				if err != nil {
					wt.abort(err)
				} else {
					wt.terminate()
				}
			} else {
				wt.abort(fmt.Errorf("server does not support write requests"))
			}
			s.wg.Done()
		}()
	case pRRQ:
		filename, mode, opts, err := unpackRQ(p)
		if err != nil {
			return fmt.Errorf("unpack RRQ: %v", err)
		}
		//fmt.Printf("got RRQ (filename=%s, mode=%s, opts=%v)\n", filename, mode, opts)
		conn, err := net.ListenUDP("udp", &net.UDPAddr{})
		if err != nil {
			return err
		}
		rf := &sender{
			send:        make([]byte, datagramLength),
			sendA:       senderAnticipate{enabled: false},
			receive:     make([]byte, datagramLength),
			tid:         remoteAddr.Port,
			conn:        conn,
			retry:       &backoff{handler: s.backoff},
			timeout:     s.timeout,
			retries:     s.retries,
			addr:        remoteAddr,
			localIP:     localAddr,
			mode:        mode,
			opts:        opts,
			maxBlockLen: maxBlockLen,
		}
		if s.sendAEnable { /* senderAnticipate if enabled in server */
			rf.sendA.enabled = true /* pass enable from server to sender */
			sendAInit(&rf.sendA, datagramLength, s.sendAWinSz)
		}
		s.wg.Add(1)
		go func() {
			if s.readHandler != nil {
				err := s.readHandler(filename, rf)
				if err != nil {
					rf.abort(err)
				}
			} else {
				rf.abort(fmt.Errorf("server does not support read requests"))
			}
			s.wg.Done()
		}()
	default:
		return fmt.Errorf("unexpected %T", p)
	}
	return nil
}
<commit_msg>Add access to remoteAddr in handlers<commit_after>package tftp

import (
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

// NewServer creates TFTP server. It requires two functions to handle
// read and write requests.
// In case nil is provided for read or write handler the respective
// operation is disabled.
func NewServer(readHandler func(filename string, rf io.ReaderFrom) error,
	writeHandler func(filename string, wt io.WriterTo) error) *Server {
	return &Server{
		readHandler:  readHandler,
		writeHandler: writeHandler,
		timeout:      defaultTimeout,
		retries:      defaultRetries,
	}
}

// NewServerWithAddr creates a new server with read and write handlers
// These handlers will have access to the remote address of the incoming connection
func NewServerWithAddr(readHandlerWithAddr func(filename string, addr *net.UDPAddr, rf io.ReaderFrom) error,
	writeWithAddrHandler func(filename string, addr *net.UDPAddr, wt io.WriterTo) error) *Server {
	return &Server{
		readWithAddrHandler:  readHandlerWithAddr,
		writeWithAddrHandler: writeWithAddrHandler,
		timeout:              defaultTimeout,
		retries:              defaultRetries,
	}
}

// RequestPacketInfo provides a method of getting the local IP address
// that is handling a UDP request.  It relies for its accuracy on the
// OS providing methods to inspect the underlying UDP and IP packets
// directly.
type RequestPacketInfo interface {
	// LocalAddr returns the IP address we are servicing the request on.
	// If it is unable to determine what address that is, the returned
	// net.IP will be nil.
	LocalIP() net.IP
}

type Server struct {
<<<<<<< HEAD
	readHandler  func(filename string, rf io.ReaderFrom) error
	writeHandler func(filename string, wt io.WriterTo) error
	backoff      backoffFunc
	conn         *net.UDPConn
	quit         chan chan struct{}
	wg           sync.WaitGroup
	timeout      time.Duration
	retries      int
	maxBlockLen  int
	sendAEnable  bool /* senderAnticipate enable by server */
	sendAWinSz   uint
=======
	readHandler          func(filename string, rf io.ReaderFrom) error
	readWithAddrHandler  func(filename string, addr *net.UDPAddr, rf io.ReaderFrom) error
	writeHandler         func(filename string, wt io.WriterTo) error
	writeWithAddrHandler func(filename string, addr *net.UDPAddr, wt io.WriterTo) error
	backoff              backoffFunc
	conn                 *net.UDPConn
	quit                 chan chan struct{}
	wg                   sync.WaitGroup
	timeout              time.Duration
	retries              int
	sendAEnable          bool /* senderAnticipate enable by server */
	sendAWinSz           uint
>>>>>>> 992b692... Add access to remoteAddr in handlers
}

// SetAnticipate provides an experimental feature in which when a packets
// is requested the server will keep sending a number of packets before
// checking whether an ack has been received. It improves tftp downloading
// speed by a few times.
// The argument winsz specifies how many packets will be sent before
// waiting for an ack packet.
// When winsz is bigger than 1, the feature is enabled, and the server
// runs through a different experimental code path. When winsz is 0 or 1,
// the feature is disabled.
func (s *Server) SetAnticipate(winsz uint) {
	if winsz > 1 {
		s.sendAEnable = true
		s.sendAWinSz = winsz
	} else {
		s.sendAEnable = false
		s.sendAWinSz = 1
	}
}

// SetTimeout sets maximum time server waits for single network
// round-trip to succeed.
// Default is 5 seconds.
func (s *Server) SetTimeout(t time.Duration) {
	if t <= 0 {
		s.timeout = defaultTimeout
	} else {
		s.timeout = t
	}
}

// SetBlockSize sets the maximum size of an individual data block.
// This must be a value between 512 (the default block size for TFTP)
// and 65456 (the max size a UDP packet payload can be).
//
// This is an advisory value -- it will be clamped to the smaller of
// the block size the client wants and the MTU of the interface being
// communicated over munis overhead.
func (s *Server) SetBlockSize(i int) {
	if i > 512 && i < 65465 {
		s.maxBlockLen = i
	}
}

// SetRetries sets maximum number of attempts server made to transmit a
// packet.
// Default is 5 attempts.
func (s *Server) SetRetries(count int) {
	if count < 1 {
		s.retries = defaultRetries
	} else {
		s.retries = count
	}
}

// SetBackoff sets a user provided function that is called to provide a
// backoff duration prior to retransmitting an unacknowledged packet.
func (s *Server) SetBackoff(h backoffFunc) {
	s.backoff = h
}

// ListenAndServe binds to address provided and start the server.
// ListenAndServe returns when Shutdown is called.
func (s *Server) ListenAndServe(addr string) error {
	a, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", a)
	if err != nil {
		return err
	}
	return s.Serve(conn)
}

// Serve starts server provided already opened UDP connecton. It is
// useful for the case when you want to run server in separate goroutine
// but still want to be able to handle any errors opening connection.
// Serve returns when Shutdown is called or connection is closed.
func (s *Server) Serve(conn *net.UDPConn) error {
	defer conn.Close()
	laddr := conn.LocalAddr()
	host, _, err := net.SplitHostPort(laddr.String())
	if err != nil {
		return err
	}
	s.conn = conn
	// Having seperate control paths for IP4 and IP6 is annoying,
	// but necessary at this point.
	addr := net.ParseIP(host)
	if addr == nil {
		return fmt.Errorf("Failed to determine IP class of listening address")
	}
	var conn4 *ipv4.PacketConn
	var conn6 *ipv6.PacketConn
	if addr.To4() != nil {
		conn4 = ipv4.NewPacketConn(conn)
		if err := conn4.SetControlMessage(ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
			conn4 = nil
		}
	} else {
		conn6 = ipv6.NewPacketConn(conn)
		if err := conn6.SetControlMessage(ipv6.FlagDst|ipv6.FlagInterface, true); err != nil {
			conn6 = nil
		}
	}

	s.quit = make(chan chan struct{})
	for {
		select {
		case q := <-s.quit:
			q <- struct{}{}
			return nil
		default:
			var err error
			if conn4 != nil {
				err = s.processRequest4(conn4)
			} else if conn6 != nil {
				err = s.processRequest6(conn6)
			} else {
				err = s.processRequest()
			}
			if err != nil {
				// TODO: add logging handler
			}
		}
	}
}

// Yes, I don't really like having seperate IPv4 and IPv6 variants,
// bit we are relying on the low-level packet control channel info to
// get a reliable source address, and those have different types and
// the struct itself is not easily interface-ized or embedded.
//
// If control is nil for whatever reason (either things not being
// implemented on a target OS or whatever other reason), localIP
// (and hence LocalIP()) will return a nil IP address.
func (s *Server) processRequest4(conn4 *ipv4.PacketConn) error {
	buf := make([]byte, datagramLength)
	cnt, control, srcAddr, err := conn4.ReadFrom(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	maxSz := blockLength
	var localAddr net.IP
	if control != nil {
		localAddr = control.Dst
		if intf, err := net.InterfaceByIndex(control.IfIndex); err == nil {
			// mtu - ipv4 overhead - udp overhead
			maxSz = intf.MTU - 28
		}
	}
	return s.handlePacket(localAddr, srcAddr.(*net.UDPAddr), buf, cnt, maxSz)
}

func (s *Server) processRequest6(conn6 *ipv6.PacketConn) error {
	buf := make([]byte, datagramLength)
	cnt, control, srcAddr, err := conn6.ReadFrom(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	maxSz := blockLength
	var localAddr net.IP
	if control != nil {
		localAddr = control.Dst
		if intf, err := net.InterfaceByIndex(control.IfIndex); err == nil {
			// mtu - ipv6 overhead - udp overhead
			maxSz = intf.MTU - 48
		}
	}
	return s.handlePacket(localAddr, srcAddr.(*net.UDPAddr), buf, cnt, maxSz)
}

// Fallback if we had problems opening a ipv4/6 control channel
func (s *Server) processRequest() error {
	buf := make([]byte, datagramLength)
	cnt, srcAddr, err := s.conn.ReadFromUDP(buf)
	if err != nil {
		return fmt.Errorf("reading UDP: %v", err)
	}
	return s.handlePacket(nil, srcAddr, buf, cnt, blockLength)
}

// Shutdown make server stop listening for new requests, allows
// server to finish outstanding transfers and stops server.
func (s *Server) Shutdown() {
	s.conn.Close()
	q := make(chan struct{})
	s.quit <- q
	<-q
	s.wg.Wait()
}

func (s *Server) handlePacket(localAddr net.IP, remoteAddr *net.UDPAddr, buffer []byte, n, maxBlockLen int) error {
	if s.maxBlockLen > 0 && s.maxBlockLen < maxBlockLen {
		maxBlockLen = s.maxBlockLen
	}
	if maxBlockLen < blockLength {
		maxBlockLen = blockLength
	}
	p, err := parsePacket(buffer[:n])
	if err != nil {
		return err
	}
	switch p := p.(type) {
	case pWRQ:
		filename, mode, opts, err := unpackRQ(p)
		if err != nil {
			return fmt.Errorf("unpack WRQ: %v", err)
		}
		//fmt.Printf("got WRQ (filename=%s, mode=%s, opts=%v)\n", filename, mode, opts)
		conn, err := net.ListenUDP("udp", &net.UDPAddr{})
		if err != nil {
			return err
		}
		if err != nil {
			return fmt.Errorf("open transmission: %v", err)
		}
		wt := &receiver{
			send:        make([]byte, datagramLength),
			receive:     make([]byte, datagramLength),
			conn:        conn,
			retry:       &backoff{handler: s.backoff},
			timeout:     s.timeout,
			retries:     s.retries,
			addr:        remoteAddr,
			localIP:     localAddr,
			mode:        mode,
			opts:        opts,
			maxBlockLen: maxBlockLen,
		}
		s.wg.Add(1)
		go func() {
			if s.writeWithAddrHandler != nil {
				err := s.writeWithAddrHandler(filename, remoteAddr, wt)
				if err != nil {
					wt.abort(err)
				} else {
					wt.terminate()
				}

			} else if s.writeHandler != nil {
				err := s.writeHandler(filename, wt)
				if err != nil {
					wt.abort(err)
				} else {
					wt.terminate()
				}
			} else {
				wt.abort(fmt.Errorf("server does not support write requests"))
			}
			s.wg.Done()
		}()
	case pRRQ:
		filename, mode, opts, err := unpackRQ(p)
		if err != nil {
			return fmt.Errorf("unpack RRQ: %v", err)
		}
		//fmt.Printf("got RRQ (filename=%s, mode=%s, opts=%v)\n", filename, mode, opts)
		conn, err := net.ListenUDP("udp", &net.UDPAddr{})
		if err != nil {
			return err
		}
		rf := &sender{
			send:        make([]byte, datagramLength),
			sendA:       senderAnticipate{enabled: false},
			receive:     make([]byte, datagramLength),
			tid:         remoteAddr.Port,
			conn:        conn,
			retry:       &backoff{handler: s.backoff},
			timeout:     s.timeout,
			retries:     s.retries,
			addr:        remoteAddr,
			localIP:     localAddr,
			mode:        mode,
			opts:        opts,
			maxBlockLen: maxBlockLen,
		}
		if s.sendAEnable { /* senderAnticipate if enabled in server */
			rf.sendA.enabled = true /* pass enable from server to sender */
			sendAInit(&rf.sendA, datagramLength, s.sendAWinSz)
		}
		s.wg.Add(1)
		go func() {
			if s.readWithAddrHandler != nil {
				err := s.readWithAddrHandler(filename, remoteAddr, rf)
				if err != nil {
					rf.abort(err)
				}
			} else if s.readHandler != nil {
				err := s.readHandler(filename, rf)
				if err != nil {
					rf.abort(err)
				}
			} else {
				rf.abort(fmt.Errorf("server does not support read requests"))
			}
			s.wg.Done()
		}()
	default:
		return fmt.Errorf("unexpected %T", p)
	}
	return nil
}
<|endoftext|>