<commit_before>package mdns

import (
	"context"
	"fmt"
	"log"
	"math/big"
	"net"
	"sync"
	"time"

	"golang.org/x/net/dns/dnsmessage"
	"golang.org/x/net/ipv4"
)

// Conn represents a mDNS Server
type Conn struct {
	mu sync.RWMutex

	socket  *ipv4.PacketConn
	dstAddr *net.UDPAddr

	queryInterval time.Duration
	localNames    []string
	queries       map[string]chan queryResult
}

type queryResult struct {
	answer dnsmessage.ResourceHeader
	addr   net.Addr
}

const (
	inboundBufferSize    = 512
	defaultQueryInterval = time.Second
	destinationAddress   = "224.0.0.251:5353"
)

// Server establishes a mDNS connection over an existing conn
func Server(conn *ipv4.PacketConn, config *Config) (*Conn, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	joinErrCount := 0
	for i := range ifaces {
		if err = conn.JoinGroup(&ifaces[i], &net.UDPAddr{IP: net.IPv4(224, 0, 0, 251)}); err != nil {
			joinErrCount++
		}
	}
	if joinErrCount >= len(ifaces) {
		return nil, errJoiningMulticastGroup
	}

	dstAddr, err := net.ResolveUDPAddr("udp", destinationAddress)
	if err != nil {
		return nil, err

	}

	c := &Conn{
		queryInterval: defaultQueryInterval,
		queries:       map[string]chan queryResult{},
		socket:        conn,
		dstAddr:       dstAddr,
	}
	if config != nil {
		if config.QueryInterval != 0 {
			c.queryInterval = config.QueryInterval
		}
		c.localNames = append([]string(nil), config.LocalNames...)
	}

	go c.start()
	return c, nil
}

func (c *Conn) sendQuestion(name string) {
	packedName, err := dnsmessage.NewName(name)
	if err != nil {
		log.Fatal(err)
	}

	msg := dnsmessage.Message{
		Header: dnsmessage.Header{},
		Questions: []dnsmessage.Question{
			{
				Type:  dnsmessage.TypeA,
				Class: dnsmessage.ClassINET,
				Name:  packedName,
			},
			{
				Type:  dnsmessage.TypeAAAA,
				Class: dnsmessage.ClassINET,
				Name:  packedName,
			},
		},
	}

	rawQuery, err := msg.Pack()
	if err != nil {
		log.Fatal(err)
	}

	if _, err := c.socket.WriteTo(rawQuery, nil, c.dstAddr); err != nil {
		log.Fatal(err)
	}
}

func ipToBytes(ip net.IP) (out [4]byte) {
	rawIP := ip.To4()
	if rawIP == nil {
		return
	}

	ipInt := big.NewInt(0)
	ipInt.SetBytes(rawIP)
	copy(out[:], ipInt.Bytes())
	return
}

func (c *Conn) sendAnswer(name string) {
	packedName, err := dnsmessage.NewName(name)
	if err != nil {
		log.Fatal(err)
	}

	msg := dnsmessage.Message{
		Header: dnsmessage.Header{
			Response:      true,
			Authoritative: true,
		},
		Answers: []dnsmessage.Resource{
			{
				Header: dnsmessage.ResourceHeader{
					Type:  dnsmessage.TypeA,
					Class: dnsmessage.ClassINET,
					Name:  packedName,
				},
				Body: &dnsmessage.AResource{
					A: ipToBytes(net.ParseIP("192.168.0.1")), // TODO actual IP
				},
			},
		},
	}

	rawAnswer, err := msg.Pack()
	if err != nil {
		log.Fatal(err)
	}

	if _, err := c.socket.WriteTo(rawAnswer, nil, c.dstAddr); err != nil {
		log.Fatal(err)
	}
}

// Query sends mDNS Queries for the following name until
// either the Context is canceled/expires or we get a result
func (c *Conn) Query(ctx context.Context, name string) (dnsmessage.ResourceHeader, net.Addr) {
	queryChan := make(chan queryResult, 1)
	c.mu.Lock()
	c.queries[name] = queryChan
	ticker := time.NewTicker(c.queryInterval)
	c.mu.Unlock()

	c.sendQuestion(name)
	for {
		select {
		case <-ticker.C:
			c.sendQuestion(name)
		case res, ok := <-queryChan:
			if !ok {
				return dnsmessage.ResourceHeader{}, nil
			}
			return res.answer, res.addr
		case <-ctx.Done():
			return dnsmessage.ResourceHeader{}, nil
		}
	}
}

func (c *Conn) start() {
	b := make([]byte, inboundBufferSize)
	p := dnsmessage.Parser{}

	for {
		n, _, src, err := c.socket.ReadFrom(b)
		if err != nil {
			log.Fatal("Read failed:", err)
			// TODO cleanup
		}

		func() {
			c.mu.RLock()
			defer c.mu.RUnlock()

			if _, err := p.Start(b[:n]); err != nil {
				fmt.Println(err)
			}

			for {
				q, err := p.Question()
				if err == dnsmessage.ErrSectionDone {
					break
				} else if err != nil {
					fmt.Println(err)
					return
				}

				for _, localName := range c.localNames {
					if localName == q.Name.String() {
						c.sendAnswer(q.Name.String())
					}
				}
			}

			for {
				a, err := p.AnswerHeader()
				if err == dnsmessage.ErrSectionDone {
					break
				}
				if err != nil {
					fmt.Println(err)
					return
				}

				if a.Type != dnsmessage.TypeA {
					continue
				}
				if resChan, ok := c.queries[a.Name.String()]; ok {
					resChan <- queryResult{a, src}
					delete(c.queries, a.Name.String())
				}

			}
		}()
	}
}
<commit_msg>We work with Chromium!<commit_after>package mdns

import (
	"context"
	"fmt"
	"log"
	"math/big"
	"net"
	"sync"
	"time"

	"golang.org/x/net/dns/dnsmessage"
	"golang.org/x/net/ipv4"
)

// Conn represents a mDNS Server
type Conn struct {
	mu sync.RWMutex

	socket  *ipv4.PacketConn
	dstAddr *net.UDPAddr

	queryInterval time.Duration
	localNames    []string
	queries       map[string]chan queryResult
}

type queryResult struct {
	answer dnsmessage.ResourceHeader
	addr   net.Addr
}

const (
	inboundBufferSize    = 512
	defaultQueryInterval = time.Second
	destinationAddress   = "224.0.0.251:5353"
	maxMessageRecords    = 3
)

// Server establishes a mDNS connection over an existing conn
func Server(conn *ipv4.PacketConn, config *Config) (*Conn, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	joinErrCount := 0
	for i := range ifaces {
		if err = conn.JoinGroup(&ifaces[i], &net.UDPAddr{IP: net.IPv4(224, 0, 0, 251)}); err != nil {
			joinErrCount++
		}
	}
	if joinErrCount >= len(ifaces) {
		return nil, errJoiningMulticastGroup
	}

	dstAddr, err := net.ResolveUDPAddr("udp", destinationAddress)
	if err != nil {
		return nil, err

	}

	c := &Conn{
		queryInterval: defaultQueryInterval,
		queries:       map[string]chan queryResult{},
		socket:        conn,
		dstAddr:       dstAddr,
	}
	if config != nil {
		if config.QueryInterval != 0 {
			c.queryInterval = config.QueryInterval
		}
		c.localNames = append([]string(nil), config.LocalNames...)
	}

	go c.start()
	return c, nil
}

func ipToBytes(ip net.IP) (out [4]byte) {
	rawIP := ip.To4()
	if rawIP == nil {
		return
	}

	ipInt := big.NewInt(0)
	ipInt.SetBytes(rawIP)
	copy(out[:], ipInt.Bytes())
	return
}

func interfaceForRemote(remote string) net.IP {
	conn, err := net.Dial("udp", remote)
	if err != nil {
		log.Fatal(err)
	}

	localAddr := conn.LocalAddr().(*net.UDPAddr)
	if err := conn.Close(); err != nil {
		log.Fatal(err)
	}

	return localAddr.IP
}

func (c *Conn) sendQuestion(name string) {
	packedName, err := dnsmessage.NewName(name)
	if err != nil {
		log.Fatal(err)
	}

	msg := dnsmessage.Message{
		Header: dnsmessage.Header{},
		Questions: []dnsmessage.Question{
			{
				Type:  dnsmessage.TypeA,
				Class: dnsmessage.ClassINET,
				Name:  packedName,
			},
		},
	}

	rawQuery, err := msg.Pack()
	if err != nil {
		log.Fatal(err)
	}

	if _, err := c.socket.WriteTo(rawQuery, nil, c.dstAddr); err != nil {
		log.Fatal(err)
	}
}

func (c *Conn) sendAnswer(name string, dst net.IP) {
	packedName, err := dnsmessage.NewName(name)
	if err != nil {
		log.Fatal(err)
	}

	msg := dnsmessage.Message{
		Header: dnsmessage.Header{
			Response:      true,
			Authoritative: true,
		},
		Answers: []dnsmessage.Resource{
			{
				Header: dnsmessage.ResourceHeader{
					Type:  dnsmessage.TypeA,
					Class: dnsmessage.ClassINET,
					Name:  packedName,
				},
				Body: &dnsmessage.AResource{
					A: ipToBytes(dst),
				},
			},
		},
	}

	rawAnswer, err := msg.Pack()
	if err != nil {
		log.Fatal(err)
	}

	if _, err := c.socket.WriteTo(rawAnswer, nil, c.dstAddr); err != nil {
		log.Fatal(err)
	}
}

// Query sends mDNS Queries for the following name until
// either the Context is canceled/expires or we get a result
func (c *Conn) Query(ctx context.Context, name string) (dnsmessage.ResourceHeader, net.Addr) {
	queryChan := make(chan queryResult, 1)
	c.mu.Lock()
	c.queries[name] = queryChan
	ticker := time.NewTicker(c.queryInterval)
	c.mu.Unlock()

	c.sendQuestion(name)
	for {
		select {
		case <-ticker.C:
			c.sendQuestion(name)
		case res, ok := <-queryChan:
			if !ok {
				return dnsmessage.ResourceHeader{}, nil
			}
			return res.answer, res.addr
		case <-ctx.Done():
			return dnsmessage.ResourceHeader{}, nil
		}
	}
}

func (c *Conn) start() {
	b := make([]byte, inboundBufferSize)
	p := dnsmessage.Parser{}

	for {
		n, _, src, err := c.socket.ReadFrom(b)
		if err != nil {
			log.Fatal("Read failed:", err)
			// TODO cleanup
		}

		func() {
			c.mu.RLock()
			defer c.mu.RUnlock()

			if _, err := p.Start(b[:n]); err != nil {
				fmt.Println(err)
			}

			for i := 0; i <= maxMessageRecords; i++ {
				q, err := p.Question()
				if err == dnsmessage.ErrSectionDone {
					break
				} else if err != nil {
					fmt.Println(err)
					return
				}

				for _, localName := range c.localNames {
					if localName == q.Name.String() {
						c.sendAnswer(q.Name.String(), interfaceForRemote(src.String()))
					}
				}
			}

			for i := 0; i <= maxMessageRecords; i++ {
				a, err := p.AnswerHeader()
				if err == dnsmessage.ErrSectionDone {
					return
				}
				if err != nil {
					fmt.Println(err)
					return
				}

				if a.Type != dnsmessage.TypeA && a.Type != dnsmessage.TypeAAAA {
					continue
				}

				if resChan, ok := c.queries[a.Name.String()]; ok {
					resChan <- queryResult{a, src}
					delete(c.queries, a.Name.String())
				}

			}
		}()
	}
}
<|endoftext|>


<commit_before>package mokumokuren

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"log"
	"net"
	"strconv"
	"sync"
	"time"
)

const CQLEN = 64

type PacketEvent struct {
	Packet    gopacket.Packet
	Timestamp time.Time
	Reverse   bool
	Finish    bool
}

type FlowKey struct {
	l3, l4 gopacket.Flow
	proto  layers.IPProtocol
}

func (key FlowKey) String() string {
	// FIXME make this prettier
	return key.l3.String() + " " + key.l4.String() + " " + strconv.Itoa(int(key.proto))
}

func NewFlowKey(sip net.IP, dip net.IP, sp uint16, dp uint16, proto layers.IPProtocol) (k FlowKey) {
	k.l3, _ = gopacket.FlowFromEndpoints(layers.NewIPEndpoint(sip), layers.NewIPEndpoint(dip))
	switch proto {
	case layers.IPProtocolTCP:
		k.l4, _ = gopacket.FlowFromEndpoints(layers.NewTCPPortEndpoint(layers.TCPPort(sp)), layers.NewTCPPortEndpoint(layers.TCPPort(dp)))
	case layers.IPProtocolUDP:
		k.l4, _ = gopacket.FlowFromEndpoints(layers.NewUDPPortEndpoint(layers.UDPPort(sp)), layers.NewUDPPortEndpoint(layers.UDPPort(dp)))
	}
	k.proto = proto
	return
}

func (k FlowKey) Reverse() FlowKey {
	return FlowKey{k.l3.Reverse(), k.l4.Reverse(), k.proto}
}

func ExtractFlowKey(pkt gopacket.Packet) (k FlowKey) {
	n := pkt.NetworkLayer()
	if n != nil {
		k.l3 = n.NetworkFlow()
		t := pkt.TransportLayer()
		if t != nil {
			k.l4 = t.TransportFlow()
		}
		switch n.(type) {
		case *layers.IPv4:
			k.proto = n.(*layers.IPv4).Protocol
		case *layers.IPv6:
			k.proto = n.(*layers.IPv6).NextHeader
		default:
			k.proto = 255
		}
	}
	return
}

// FIXME: consider making FlowEntry an interface
// and add a factory instead of doing runtime typing here.
// ...runtime typing is kind of pythonic...
type FlowEntry struct {
	Key FlowKey

	Counters  map[int]int
	Data      map[int]interface{}
	StartTime time.Time
	LastTime  time.Time

	packetChannel chan PacketEvent
	reapChannel   chan FlowKey
	flowFinishing chan struct{}
	flowDone      chan struct{}
}

// Mark this flow as finished. To be called by a layer chain function
// on the last packet of the flow, or by an idle flow reaper.

func (fe *FlowEntry) Finish() {
	fe.reapChannel <- fe.Key
}

func (fe *FlowEntry) String() string {
	// FIXME this gets fixed when moving to statically typed flow entries
	return fmt.Sprintf("%v", *fe)
}

type FlowChainFn func(*FlowEntry) bool

type PacketChainFn func(*FlowEntry, PacketEvent) bool

type LayerChainFn func(*FlowEntry, PacketEvent, gopacket.Layer) bool

type layerChainEntry struct {
	LayerType gopacket.LayerType
	Fn        LayerChainFn
}

const IDLE_TIMEOUT = 30

// Contains the set of currently active flows
type FlowTable struct {
	// Function chain to run on the first packet in any given flow.
	// Used to set up state for subsequent chains
	initialChain []PacketChainFn

	// Function chain to run for each layer of each incoming packet
	layerChain []layerChainEntry

	// Function chain to run on concluded flows
	emitterChain []FlowChainFn

	// The current time as of the last packet added to the flow
	packetClock time.Time

	// Currently active flows, maps a flow key to a flow entry.
	active map[FlowKey]*FlowEntry

	// Idle queue
	idleq IdleQueue

	// Lock guarding access to the active table
	activeLock sync.RWMutex

	// Channel for flow entries to be reaped from the active queue
	reapChannel chan FlowKey
	reaperDone  chan struct{}

	// Channel for packet clock ticks
	tickChannel chan time.Time
}

func NewFlowTable() (ft *FlowTable) {
	ft = new(FlowTable)
	ft.initialChain = make([]PacketChainFn, 0)
	ft.layerChain = make([]layerChainEntry, 0)
	ft.emitterChain = make([]FlowChainFn, 0)
	ft.active = make(map[FlowKey]*FlowEntry)
	ft.reapChannel = make(chan FlowKey, CQLEN)
	ft.reaperDone = make(chan struct{})
	ft.tickChannel = make(chan time.Time)

	// start the flow table's emitter
	go ft.reapFinishedFlowEntries()

	// start the idle reaper
	go ft.reapIdleFlowEntries()

	return
}

func (ft *FlowTable) HandlePacket(pkt gopacket.Packet) {
	// advance the packet clock
	timestamp := pkt.Metadata().Timestamp
	ft.tickPacketClock(timestamp)

	// extract a flow key from the packet
	k := ExtractFlowKey(pkt)

	// get a flow entry for the flow key, tick the idle queue,
	// and send it the packet for further processing if not ignored.
	fe, rev := ft.flowEntry(k)
	if fe != nil {
		ft.idleq.Tick(k, timestamp)
		fe.packetChannel <- PacketEvent{pkt, timestamp, rev, false}
	}
}

func (ft *FlowTable) Shutdown() {

	log.Printf("in shutdown")

	ft.activeLock.RLock()
	keylist := make([]FlowKey, len(ft.active))
	for k := range ft.active {
		keylist = append(keylist, k)
	}
	ft.activeLock.RUnlock()

	log.Printf("reaping %u keys", len(keylist))

	for _, k := range keylist {
		ft.reapChannel <- k
	}

	// Shut down the reapers
	close(ft.tickChannel)
	close(ft.reapChannel)

	// and wait
	_ = <-ft.reaperDone
}

func (ft *FlowTable) AddInitialFunction(fn PacketChainFn) {
	ft.initialChain = append(ft.initialChain, fn)
	log.Printf("initialChain now %v", ft.initialChain)
}

func (ft *FlowTable) AddLayerFunction(fn LayerChainFn, layerType gopacket.LayerType) {
	ft.layerChain = append(ft.layerChain, layerChainEntry{layerType, fn})
	log.Printf("layerChain now %v", ft.layerChain)
}

func (ft *FlowTable) AddEmitterFunction(fn FlowChainFn) {
	ft.emitterChain = append(ft.emitterChain, fn)
	log.Printf("emitterChain now %v", ft.emitterChain)
}

func (ft *FlowTable) flowEntry(key FlowKey) (fe *FlowEntry, rev bool) {
	// First look for a flow entry in the active table
	ft.activeLock.RLock()
	fe = ft.active[key]
	ft.activeLock.RUnlock()

	if fe != nil {
		return fe, false
	}

	// Now look for a reverse flow entry
	ft.activeLock.RLock()
	fe = ft.active[key.Reverse()]
	ft.activeLock.RUnlock()

	if fe != nil {
		return fe, true
	}

	// No entry available. Create a new one.
	fe = new(FlowEntry)
	fe.Key = key
	fe.Counters = make(map[int]int)
	fe.Data = make(map[int]interface{})
	fe.packetChannel = make(chan PacketEvent)
	fe.reapChannel = ft.reapChannel
	fe.flowFinishing = make(chan struct{})
	fe.flowDone = make(chan struct{})

	// Now start running the function chain for this flow entry
	// in its own goroutine
	go func() {
		initial := true

		// Run until we get a finishing signal.
		for {
			select {
			case pe := <-fe.packetChannel:
				if fe.LastTime.Before(pe.Timestamp) {
					fe.LastTime = pe.Timestamp
				}

				// The initial function chain is used to set up
				// state in the FlowEntry to be used by
				// functions in the layer chain.
				if initial {
					fe.StartTime = pe.Timestamp
					for _, fn := range ft.initialChain {
						if !fn(fe, pe) {
							break
						}
					}
					initial = false
				}

				// The layer chain contains functions to be
				// called for specified layers in a
				// specified order
				for _, le := range ft.layerChain {
					layer := pe.Packet.Layer(le.LayerType)
					if layer != nil {
						if !le.Fn(fe, pe, layer) {
							break
						}
					}
				}
			case <-fe.flowFinishing:
				break
			}
		}

		log.Printf("flowroutine for key %v exiting", fe.Key)

		// let the reaper know we're done
		close(fe.flowDone)
	}()

	// Add the flow to the active table
	ft.activeLock.Lock()
	ft.active[key] = fe
	ft.activeLock.Unlock()
	//log.Printf("added new flow entry for key %s", key.String())

	return fe, false
}

func quantize(tick time.Time, quantum int64) time.Time {
	return time.Unix((tick.Unix()/quantum)*quantum, 0)
}

const BIN_QUANTUM = 10

func (ft *FlowTable) tickPacketClock(tick time.Time) {
	if ft.packetClock.After(tick) {
		return
	}

	var lastClock time.Time
	if ft.packetClock.After(time.Time{}) {
		lastClock = ft.packetClock
	} else {
		lastClock = tick
		log.Printf("initializing packet clock %v", tick)
	}

	ft.packetClock = tick

	// generate ticks for reapIdleFlowEntries
	for i, j := quantize(lastClock, BIN_QUANTUM), quantize(tick, BIN_QUANTUM); i.Before(j); i = i.Add(BIN_QUANTUM * time.Second) {
		log.Printf("ticking to %v", i)
		ft.tickChannel <- i
	}

}

func (ft *FlowTable) reapIdleFlowEntries() {
	log.Println("idle reaper up")

	for tick := range ft.tickChannel {
		endtime := tick.Add(IDLE_TIMEOUT * time.Second)
		for {
			k, ok := ft.idleq.NextIdleBefore(endtime)
			if ok {
				ft.reapChannel <- k
			} else {
				break
			}
		}
	}
	log.Println("idle reaper down")
}

func (ft *FlowTable) reapFinishedFlowEntries() {
	log.Println("finished reaper up")

	for k := range ft.reapChannel {
		// get the flow
		ft.activeLock.RLock()
		fe := ft.active[k]
		ft.activeLock.RUnlock()

		// remove the flow from the active table
		ft.activeLock.Lock()
		delete(ft.active, k)
		ft.activeLock.Unlock()

		// signal flow's goroutine to complete
		fe.flowFinishing <- struct{}{}

		// and wait for it to do so
		_ = <-fe.flowDone

		// now run the emitter chain
		for _, fn := range ft.emitterChain {
			if !fn(fe) {
				break
			}
		}
	}

	// signal done
	close(ft.reaperDone)

	log.Println("finished reaper down")
}
<commit_msg>replace Endpoint/Flow with native FlowKey (build checkpoint)<commit_after>package mokumokuren

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"log"
	//"net"
	//"strconv"
	"sync"
	"time"
)

const CQLEN = 64

type PacketEvent struct {
	Packet    gopacket.Packet
	Timestamp time.Time
	Reverse   bool
	Finish    bool
}

type FlowKey struct {
	Sip string // FIXME this is a stupid hack but net.IP isn't comparable
	Dip string // FIXME this is a stupid hack but net.IP isn't comparable
	Sp  uint16
	Dp  uint16
	P   uint8
}

// Return a string representation of this FlowKey suitable for printing
func (key FlowKey) String() string {
	return fmt.Sprintf("%s:%u -> %s:%u %u", key.Sip, key.Dip, key.Sp, key.Dp, key.P)
}

// Return the reverse of this FlowKey,
// with source and destination address and port flipped.
func (k FlowKey) Reverse() FlowKey {
	return FlowKey{k.Dip, k.Sip, k.Dp, k.Sp, k.P}
}

// Extract a flow key from a packet
func ExtractFlowKey(pkt gopacket.Packet) (k FlowKey) {
	nl := pkt.NetworkLayer()
	if nl == nil {
		// ain't nobody got time for non-IP packets. empty flow key.
		return
	}

	switch nl.(type) {
	case *layers.IPv4:
		k.Sip = nl.(*layers.IPv4).SrcIP.String()
		k.Dip = nl.(*layers.IPv4).SrcIP.String()
		k.P = uint8(nl.(*layers.IPv4).Protocol)
	case *layers.IPv6:
		k.Sip = nl.(*layers.IPv6).SrcIP.String()
		k.Dip = nl.(*layers.IPv6).SrcIP.String()
		k.P = uint8(nl.(*layers.IPv6).NextHeader)
	default:
		// um i got nothing, empty flow key.
		return
	}

	tl := pkt.TransportLayer()
	if tl == nil {
		// no transport layer, so try to decode ICMP
		if icmpl := pkt.Layer(layers.LayerTypeICMPv4).(*layers.ICMPv4); icmpl != nil {
			icmptype := icmpl.TypeCode.Type()
			if icmptype == layers.ICMPv4TypeDestinationUnreachable ||
				icmptype == layers.ICMPv4TypeTimeExceeded ||
				icmptype == layers.ICMPv4TypeParameterProblem {
				// Account ICMPv4 messages from routers to the
				// reverse flow they belong to
				sk := ExtractFlowKey(gopacket.NewPacket(icmpl.LayerPayload(),
					layers.LayerTypeIPv4,
					gopacket.Default))
				k.Sip = sk.Dip
				k.Sp = sk.Dp
				k.P = sk.P
			} else {
				k.Sp = uint16(icmpl.TypeCode.Type())
				k.Dp = uint16(icmpl.TypeCode.Code())
			}
		} else if icmpl := pkt.Layer(layers.LayerTypeICMPv6).(*layers.ICMPv6); icmpl != nil {
			icmptype := icmpl.TypeCode.Type()
			if icmptype == layers.ICMPv6TypeDestinationUnreachable ||
				icmptype == layers.ICMPv6TypeTimeExceeded ||
				icmptype == layers.ICMPv6TypePacketTooBig ||
				icmptype == layers.ICMPv6TypeParameterProblem {
				// Account ICMPv6 messages from routers to the
				// reverse flow they belong to
				sk := ExtractFlowKey(gopacket.NewPacket(icmpl.LayerPayload(),
					layers.LayerTypeIPv6,
					gopacket.Default))
				k.Sip = sk.Dip
				k.Sp = sk.Dp
				k.P = sk.P
			} else {
				k.Sp = uint16(icmpl.TypeCode.Type())
				k.Dp = uint16(icmpl.TypeCode.Code())
			}
		} else {
			// no icmp, no transport, no ports for you
			return
		}
	}

	switch tl.(type) {
	case *layers.TCP:
		k.Sp = uint16(tl.(*layers.TCP).SrcPort)
		k.Dp = uint16(tl.(*layers.TCP).DstPort)
	case *layers.UDP:
		k.Sp = uint16(tl.(*layers.UDP).SrcPort)
		k.Dp = uint16(tl.(*layers.UDP).DstPort)
	case *layers.UDPLite:
		k.Sp = uint16(tl.(*layers.UDPLite).SrcPort)
		k.Dp = uint16(tl.(*layers.UDPLite).DstPort)
	case *layers.SCTP:
		k.Sp = uint16(tl.(*layers.SCTP).SrcPort)
		k.Dp = uint16(tl.(*layers.SCTP).DstPort)
	default:
		// no transport layer we know about, so no ports.
		return
	}

	// key set
	return
}

// FIXME: consider making FlowEntry an interface
// and add a factory instead of doing runtime typing here.
// ...runtime typing is kind of pythonic...
type FlowEntry struct {
	Key FlowKey

	Counters  map[int]int
	Data      map[int]interface{}
	StartTime time.Time
	LastTime  time.Time

	packetChannel chan PacketEvent
	reapChannel   chan FlowKey
	flowFinishing chan struct{}
	flowDone      chan struct{}
}

// Mark this flow as finished. To be called by a layer chain function
// on the last packet of the flow, or by an idle flow reaper.

func (fe *FlowEntry) Finish() {
	fe.reapChannel <- fe.Key
}

func (fe *FlowEntry) String() string {
	// FIXME this gets fixed when moving to statically typed flow entries
	return fmt.Sprintf("%v", *fe)
}

type FlowChainFn func(*FlowEntry) bool

type PacketChainFn func(*FlowEntry, PacketEvent) bool

type LayerChainFn func(*FlowEntry, PacketEvent, gopacket.Layer) bool

type layerChainEntry struct {
	LayerType gopacket.LayerType
	Fn        LayerChainFn
}

const IDLE_TIMEOUT = 30

// Contains the set of currently active flows
type FlowTable struct {
	// Function chain to run on the first packet in any given flow.
	// Used to set up state for subsequent chains
	initialChain []PacketChainFn

	// Function chain to run for each layer of each incoming packet
	layerChain []layerChainEntry

	// Function chain to run on concluded flows
	emitterChain []FlowChainFn

	// The current time as of the last packet added to the flow
	packetClock time.Time

	// Currently active flows, maps a flow key to a flow entry.
	active map[FlowKey]*FlowEntry

	// Idle queue
	idleq IdleQueue

	// Lock guarding access to the active table
	activeLock sync.RWMutex

	// Channel for flow entries to be reaped from the active queue
	reapChannel chan FlowKey
	reaperDone  chan struct{}

	// Channel for packet clock ticks
	tickChannel chan time.Time
}

func NewFlowTable() (ft *FlowTable) {
	ft = new(FlowTable)
	ft.initialChain = make([]PacketChainFn, 0)
	ft.layerChain = make([]layerChainEntry, 0)
	ft.emitterChain = make([]FlowChainFn, 0)
	ft.active = make(map[FlowKey]*FlowEntry)
	ft.reapChannel = make(chan FlowKey, CQLEN)
	ft.reaperDone = make(chan struct{})
	ft.tickChannel = make(chan time.Time)

	// start the flow table's emitter
	go ft.reapFinishedFlowEntries()

	// start the idle reaper
	go ft.reapIdleFlowEntries()

	return
}

func (ft *FlowTable) HandlePacket(pkt gopacket.Packet) {
	// advance the packet clock
	timestamp := pkt.Metadata().Timestamp
	ft.tickPacketClock(timestamp)

	// extract a flow key from the packet
	k := ExtractFlowKey(pkt)

	// get a flow entry for the flow key, tick the idle queue,
	// and send it the packet for further processing if not ignored.
	fe, rev := ft.flowEntry(k)
	if fe != nil {
		ft.idleq.Tick(k, timestamp)
		fe.packetChannel <- PacketEvent{pkt, timestamp, rev, false}
	}
}

func (ft *FlowTable) Shutdown() {

	log.Printf("in shutdown")

	ft.activeLock.RLock()
	keylist := make([]FlowKey, len(ft.active))
	for k := range ft.active {
		keylist = append(keylist, k)
	}
	ft.activeLock.RUnlock()

	for _, k := range keylist {
		ft.reapChannel <- k
		log.Printf("reaped %v", k)

	}

	// Shut down the reapers
	close(ft.tickChannel)
	close(ft.reapChannel)

	// and wait
	_ = <-ft.reaperDone
}

func (ft *FlowTable) AddInitialFunction(fn PacketChainFn) {
	ft.initialChain = append(ft.initialChain, fn)
	log.Printf("initialChain now %v", ft.initialChain)
}

func (ft *FlowTable) AddLayerFunction(fn LayerChainFn, layerType gopacket.LayerType) {
	ft.layerChain = append(ft.layerChain, layerChainEntry{layerType, fn})
	log.Printf("layerChain now %v", ft.layerChain)
}

func (ft *FlowTable) AddEmitterFunction(fn FlowChainFn) {
	ft.emitterChain = append(ft.emitterChain, fn)
	log.Printf("emitterChain now %v", ft.emitterChain)
}

func (ft *FlowTable) flowEntry(key FlowKey) (fe *FlowEntry, rev bool) {
	// First look for a flow entry in the active table
	ft.activeLock.RLock()
	fe = ft.active[key]
	ft.activeLock.RUnlock()

	if fe != nil {
		return fe, false
	}

	// Now look for a reverse flow entry
	ft.activeLock.RLock()
	fe = ft.active[key.Reverse()]
	ft.activeLock.RUnlock()

	if fe != nil {
		return fe, true
	}

	// No entry available. Create a new one.
	fe = new(FlowEntry)
	fe.Key = key
	fe.Counters = make(map[int]int)
	fe.Data = make(map[int]interface{})
	fe.packetChannel = make(chan PacketEvent)
	fe.reapChannel = ft.reapChannel
	fe.flowFinishing = make(chan struct{})
	fe.flowDone = make(chan struct{})

	// Now start running the function chain for this flow entry
	// in its own goroutine
	go func() {
		initial := true

		// Run until we get a finishing signal.
		for {
			select {
			case pe := <-fe.packetChannel:
				if fe.LastTime.Before(pe.Timestamp) {
					fe.LastTime = pe.Timestamp
				}

				// The initial function chain is used to set up
				// state in the FlowEntry to be used by
				// functions in the layer chain.
				if initial {
					fe.StartTime = pe.Timestamp
					for _, fn := range ft.initialChain {
						if !fn(fe, pe) {
							break
						}
					}
					initial = false
				}

				// The layer chain contains functions to be
				// called for specified layers in a
				// specified order
				for _, le := range ft.layerChain {
					layer := pe.Packet.Layer(le.LayerType)
					if layer != nil {
						if !le.Fn(fe, pe, layer) {
							break
						}
					}
				}
			case <-fe.flowFinishing:
				break
			}
		}

		log.Printf("flowroutine for key %v exiting", fe.Key)

		// let the reaper know we're done
		close(fe.flowDone)
	}()

	// Add the flow to the active table
	ft.activeLock.Lock()
	ft.active[key] = fe
	ft.activeLock.Unlock()
	//log.Printf("added new flow entry for key %s", key.String())

	return fe, false
}

func quantize(tick time.Time, quantum int64) time.Time {
	return time.Unix((tick.Unix()/quantum)*quantum, 0)
}

const BIN_QUANTUM = 10

func (ft *FlowTable) tickPacketClock(tick time.Time) {
	if ft.packetClock.After(tick) {
		return
	}

	var lastClock time.Time
	if ft.packetClock.After(time.Time{}) {
		lastClock = ft.packetClock
	} else {
		lastClock = tick
		log.Printf("initializing packet clock %v", tick)
	}

	ft.packetClock = tick

	// generate ticks for reapIdleFlowEntries
	for i, j := quantize(lastClock, BIN_QUANTUM), quantize(tick, BIN_QUANTUM); i.Before(j); i = i.Add(BIN_QUANTUM * time.Second) {
		log.Printf("ticking to %v", i)
		ft.tickChannel <- i
	}

}

func (ft *FlowTable) reapIdleFlowEntries() {
	log.Println("idle reaper up")

	for tick := range ft.tickChannel {
		endtime := tick.Add(IDLE_TIMEOUT * time.Second)
		for {
			k, ok := ft.idleq.NextIdleBefore(endtime)
			if ok {
				ft.reapChannel <- k
			} else {
				break
			}
		}
	}
	log.Println("idle reaper down")
}

func (ft *FlowTable) reapFinishedFlowEntries() {
	log.Println("finished reaper up")

	for k := range ft.reapChannel {
		// get the flow
		ft.activeLock.RLock()
		fe := ft.active[k]
		ft.activeLock.RUnlock()

		// remove the flow from the active table
		ft.activeLock.Lock()
		delete(ft.active, k)
		ft.activeLock.Unlock()

		// signal flow's goroutine to complete
		fe.flowFinishing <- struct{}{}

		// and wait for it to do so
		_ = <-fe.flowDone

		// now run the emitter chain
		for _, fn := range ft.emitterChain {
			if !fn(fe) {
				break
			}
		}
	}

	// signal done
	close(ft.reaperDone)

	log.Println("finished reaper down")
}
<|endoftext|>


<commit_before>package cos

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/hashicorp/terraform/internal/backend"
	"github.com/hashicorp/terraform/internal/legacy/helper/schema"
	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
	tag "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/tag/v20180813"
	"github.com/tencentyun/cos-go-sdk-v5"
)

// Default value from environment variable
const (
	PROVIDER_SECRET_ID  = "TENCENTCLOUD_SECRET_ID"
	PROVIDER_SECRET_KEY = "TENCENTCLOUD_SECRET_KEY"
	PROVIDER_REGION     = "TENCENTCLOUD_REGION"
)

// Backend implements "backend".Backend for tencentCloud cos
type Backend struct {
	*schema.Backend

	cosContext context.Context
	cosClient  *cos.Client
	tagClient  *tag.Client

	region  string
	bucket  string
	prefix  string
	key     string
	encrypt bool
	acl     string
}

// New creates a new backend for TencentCloud cos remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"secret_id": {
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc(PROVIDER_SECRET_ID, nil),
				Description: "Secret id of Tencent Cloud",
			},
			"secret_key": {
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc(PROVIDER_SECRET_KEY, nil),
				Description: "Secret key of Tencent Cloud",
				Sensitive:   true,
			},
			"region": {
				Type:         schema.TypeString,
				Required:     true,
				DefaultFunc:  schema.EnvDefaultFunc(PROVIDER_REGION, nil),
				Description:  "The region of the COS bucket",
				InputDefault: "ap-guangzhou",
			},
			"bucket": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the COS bucket",
			},
			"prefix": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The directory for saving the state file in bucket",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					prefix := v.(string)
					if strings.HasPrefix(prefix, "/") || strings.HasPrefix(prefix, "./") {
						return nil, []error{fmt.Errorf("prefix must not start with '/' or './'")}
					}
					return nil, nil
				},
			},
			"key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The path for saving the state file in bucket",
				Default:     "terraform.tfstate",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					if strings.HasPrefix(v.(string), "/") || strings.HasSuffix(v.(string), "/") {
						return nil, []error{fmt.Errorf("key can not start and end with '/'")}
					}
					return nil, nil
				},
			},
			"encrypt": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Whether to enable server side encryption of the state file",
				Default:     true,
			},
			"acl": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Object ACL to be applied to the state file",
				Default:     "private",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					value := v.(string)
					if value != "private" && value != "public-read" {
						return nil, []error{fmt.Errorf(
							"acl value invalid, expected %s or %s, got %s",
							"private", "public-read", value)}
					}
					return nil, nil
				},
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure

	return result
}

// configure init cos client
func (b *Backend) configure(ctx context.Context) error {
	if b.cosClient != nil {
		return nil
	}

	b.cosContext = ctx
	data := schema.FromContextBackendConfig(b.cosContext)

	b.region = data.Get("region").(string)
	b.bucket = data.Get("bucket").(string)
	b.prefix = data.Get("prefix").(string)
	b.key = data.Get("key").(string)
	b.encrypt = data.Get("encrypt").(bool)
	b.acl = data.Get("acl").(string)

	u, err := url.Parse(fmt.Sprintf("https://%s.cos.%s.myqcloud.com", b.bucket, b.region))
	if err != nil {
		return err
	}

	b.cosClient = cos.NewClient(
		&cos.BaseURL{BucketURL: u},
		&http.Client{
			Timeout: 60 * time.Second,
			Transport: &cos.AuthorizationTransport{
				SecretID:  data.Get("secret_id").(string),
				SecretKey: data.Get("secret_key").(string),
			},
		},
	)

	credential := common.NewCredential(
		data.Get("secret_id").(string),
		data.Get("secret_key").(string),
	)

	cpf := profile.NewClientProfile()
	cpf.HttpProfile.ReqMethod = "POST"
	cpf.HttpProfile.ReqTimeout = 300
	cpf.Language = "en-US"
	b.tagClient, err = tag.NewClient(credential, b.region, cpf)

	return err
}
<commit_msg>cos backend support accelerate (#31425)<commit_after>package cos

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/hashicorp/terraform/internal/backend"
	"github.com/hashicorp/terraform/internal/legacy/helper/schema"
	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common"
	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
	tag "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/tag/v20180813"
	"github.com/tencentyun/cos-go-sdk-v5"
)

// Default value from environment variable
const (
	PROVIDER_SECRET_ID  = "TENCENTCLOUD_SECRET_ID"
	PROVIDER_SECRET_KEY = "TENCENTCLOUD_SECRET_KEY"
	PROVIDER_REGION     = "TENCENTCLOUD_REGION"
)

// Backend implements "backend".Backend for tencentCloud cos
type Backend struct {
	*schema.Backend

	cosContext context.Context
	cosClient  *cos.Client
	tagClient  *tag.Client

	region  string
	bucket  string
	prefix  string
	key     string
	encrypt bool
	acl     string
}

// New creates a new backend for TencentCloud cos remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"secret_id": {
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc(PROVIDER_SECRET_ID, nil),
				Description: "Secret id of Tencent Cloud",
			},
			"secret_key": {
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc(PROVIDER_SECRET_KEY, nil),
				Description: "Secret key of Tencent Cloud",
				Sensitive:   true,
			},
			"region": {
				Type:         schema.TypeString,
				Required:     true,
				DefaultFunc:  schema.EnvDefaultFunc(PROVIDER_REGION, nil),
				Description:  "The region of the COS bucket",
				InputDefault: "ap-guangzhou",
			},
			"bucket": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the COS bucket",
			},
			"prefix": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The directory for saving the state file in bucket",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					prefix := v.(string)
					if strings.HasPrefix(prefix, "/") || strings.HasPrefix(prefix, "./") {
						return nil, []error{fmt.Errorf("prefix must not start with '/' or './'")}
					}
					return nil, nil
				},
			},
			"key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The path for saving the state file in bucket",
				Default:     "terraform.tfstate",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					if strings.HasPrefix(v.(string), "/") || strings.HasSuffix(v.(string), "/") {
						return nil, []error{fmt.Errorf("key can not start and end with '/'")}
					}
					return nil, nil
				},
			},
			"encrypt": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Whether to enable server side encryption of the state file",
				Default:     true,
			},
			"acl": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Object ACL to be applied to the state file",
				Default:     "private",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					value := v.(string)
					if value != "private" && value != "public-read" {
						return nil, []error{fmt.Errorf(
							"acl value invalid, expected %s or %s, got %s",
							"private", "public-read", value)}
					}
					return nil, nil
				},
			},
			"accelerate": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Whether to enable global Acceleration",
				Default:     false,
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure

	return result
}

// configure init cos client
func (b *Backend) configure(ctx context.Context) error {
	if b.cosClient != nil {
		return nil
	}

	b.cosContext = ctx
	data := schema.FromContextBackendConfig(b.cosContext)

	b.region = data.Get("region").(string)
	b.bucket = data.Get("bucket").(string)
	b.prefix = data.Get("prefix").(string)
	b.key = data.Get("key").(string)
	b.encrypt = data.Get("encrypt").(bool)
	b.acl = data.Get("acl").(string)

	var (
		u   *url.URL
		err error
	)
	accelerate := data.Get("accelerate").(bool)
	if accelerate {
		u, err = url.Parse(fmt.Sprintf("https://%s.cos.accelerate.myqcloud.com", b.bucket))
	} else {
		u, err = url.Parse(fmt.Sprintf("https://%s.cos.%s.myqcloud.com", b.bucket, b.region))
	}
	if err != nil {
		return err
	}

	b.cosClient = cos.NewClient(
		&cos.BaseURL{BucketURL: u},
		&http.Client{
			Timeout: 60 * time.Second,
			Transport: &cos.AuthorizationTransport{
				SecretID:  data.Get("secret_id").(string),
				SecretKey: data.Get("secret_key").(string),
			},
		},
	)

	credential := common.NewCredential(
		data.Get("secret_id").(string),
		data.Get("secret_key").(string),
	)

	cpf := profile.NewClientProfile()
	cpf.HttpProfile.ReqMethod = "POST"
	cpf.HttpProfile.ReqTimeout = 300
	cpf.Language = "en-US"
	b.tagClient, err = tag.NewClient(credential, b.region, cpf)

	return err
}
<|endoftext|>


<commit_before>// Copyright 2020 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package logstream_test

import (
	"context"
	"os"
	"path/filepath"
	"sync"
	"syscall"
	"testing"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/tailer/logstream"
	"github.com/google/mtail/internal/testutil"
	"github.com/google/mtail/internal/waker"
	"golang.org/x/sys/unix"
)

func TestPipeStreamReadCompletedBecauseClosed(t *testing.T) {
	var wg sync.WaitGroup

	tmpDir, rmTmpDir := testutil.TestTempDir(t)
	defer rmTmpDir()

	name := filepath.Join(tmpDir, "fifo")
	testutil.FatalIfErr(t, unix.Mkfifo(name, 0666))

	lines := make(chan *logline.LogLine, 1)
	ctx, cancel := context.WithCancel(context.Background())
	waker, awaken := waker.NewTest(ctx, 1)

	ps, err := logstream.New(ctx, &wg, waker, name, lines, false)
	testutil.FatalIfErr(t, err)

	f, err := os.OpenFile(name, os.O_WRONLY, os.ModeNamedPipe)
	testutil.FatalIfErr(t, err)
	testutil.WriteString(t, f, "1\n")
	awaken(1)

	// Pipes need to be closed to signal to the pipeStream to finish up.
	testutil.FatalIfErr(t, f.Close())
	awaken(0)

	ps.Stop() // no-op for pipes
	wg.Wait()
	close(lines)

	received := testutil.LinesReceived(lines)
	expected := []*logline.LogLine{
		{context.TODO(), name, "1"},
	}
	testutil.ExpectNoDiff(t, expected, received, testutil.IgnoreFields(logline.LogLine{}, "Context"))

	cancel()

	if !ps.IsComplete() {
		t.Errorf("expecting pipestream to be complete because fifo closed")
	}
}

func TestPipeStreamReadCompletedBecauseCancel(t *testing.T) {
	var wg sync.WaitGroup

	tmpDir, rmTmpDir := testutil.TestTempDir(t)
	defer rmTmpDir()

	name := filepath.Join(tmpDir, "fifo")
	testutil.FatalIfErr(t, unix.Mkfifo(name, 0666))

	lines := make(chan *logline.LogLine, 1)
	ctx, cancel := context.WithCancel(context.Background())
	waker, awaken := waker.NewTest(ctx, 1)

	ps, err := logstream.New(ctx, &wg, waker, name, lines, false)
	testutil.FatalIfErr(t, err)

	f, err := os.OpenFile(name, os.O_WRONLY|syscall.O_NONBLOCK, os.ModeNamedPipe)
	testutil.FatalIfErr(t, err)
	testutil.WriteString(t, f, "1\n")
	awaken(1)

	cancel()
	wg.Wait()
	close(lines)

	received := testutil.LinesReceived(lines)
	expected := []*logline.LogLine{
		{context.TODO(), name, "1"},
	}
	testutil.ExpectNoDiff(t, expected, received, testutil.IgnoreFields(logline.LogLine{}, "Context"))

	if !ps.IsComplete() {
		t.Errorf("expecting pipestream to be complete because cancelled")
	}
}
<commit_msg>Use the always-waker in the pipestream test.<commit_after>// Copyright 2020 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package logstream_test

import (
	"context"
	"os"
	"path/filepath"
	"sync"
	"testing"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/tailer/logstream"
	"github.com/google/mtail/internal/testutil"
	"github.com/google/mtail/internal/waker"
	"golang.org/x/sys/unix"
)

func TestPipeStreamReadCompletedBecauseClosed(t *testing.T) {
	var wg sync.WaitGroup

	tmpDir, rmTmpDir := testutil.TestTempDir(t)
	defer rmTmpDir()

	name := filepath.Join(tmpDir, "fifo")
	testutil.FatalIfErr(t, unix.Mkfifo(name, 0666))

	lines := make(chan *logline.LogLine, 1)
	ctx, cancel := context.WithCancel(context.Background())
	waker := waker.NewTestAlways()

	ps, err := logstream.New(ctx, &wg, waker, name, lines, false)
	testutil.FatalIfErr(t, err)

	f, err := os.OpenFile(name, os.O_WRONLY, os.ModeNamedPipe)
	testutil.FatalIfErr(t, err)
	testutil.WriteString(t, f, "1\n")

	// Pipes need to be closed to signal to the pipeStream to finish up.
	testutil.FatalIfErr(t, f.Close())

	ps.Stop() // no-op for pipes
	wg.Wait()
	close(lines)

	received := testutil.LinesReceived(lines)
	expected := []*logline.LogLine{
		{context.TODO(), name, "1"},
	}
	testutil.ExpectNoDiff(t, expected, received, testutil.IgnoreFields(logline.LogLine{}, "Context"))

	cancel()

	if !ps.IsComplete() {
		t.Errorf("expecting pipestream to be complete because fifo closed")
	}
}

func TestPipeStreamReadCompletedBecauseCancel(t *testing.T) {
	var wg sync.WaitGroup

	tmpDir, rmTmpDir := testutil.TestTempDir(t)
	defer rmTmpDir()

	name := filepath.Join(tmpDir, "fifo")
	testutil.FatalIfErr(t, unix.Mkfifo(name, 0666))

	lines := make(chan *logline.LogLine, 1)
	ctx, cancel := context.WithCancel(context.Background())
	waker := waker.NewTestAlways()

	ps, err := logstream.New(ctx, &wg, waker, name, lines, false)
	testutil.FatalIfErr(t, err)

	f, err := os.OpenFile(name, os.O_WRONLY, os.ModeNamedPipe)
	testutil.FatalIfErr(t, err)
	testutil.WriteString(t, f, "1\n")

	cancel()
	wg.Wait()
	close(lines)

	received := testutil.LinesReceived(lines)
	expected := []*logline.LogLine{
		{context.TODO(), name, "1"},
	}
	testutil.ExpectNoDiff(t, expected, received, testutil.IgnoreFields(logline.LogLine{}, "Context"))

	if !ps.IsComplete() {
		t.Errorf("expecting pipestream to be complete because cancelled")
	}
}
<|endoftext|>


<commit_before>package gopherstack

import (
	"fmt"
	"log"
	"time"
)

// waitForAsyncJob simply blocks until the the asynchronous job has
// executed or has timed out.
func (c CloudStackClient) WaitForAsyncJob(jobId string, timeout time.Duration) error {
	done := make(chan struct{})
	defer close(done)

	result := make(chan error, 1)
	go func() {
		attempts := 0
		for {
			attempts += 1

			log.Printf("Checking async job status... (attempt: %d)", attempts)
			status, err := c.QueryAsyncJobResult(jobId)
			if err != nil {
				result <- err
				return
			}

			// check what the real state will be.
			if status == 1 {
				result <- nil
				return
			}

			// Wait 3 seconds in between
			time.Sleep(3 * time.Second)

			// Verify we shouldn't exit
			select {
			case <-done:
				// We finished, so just exit the goroutine
				return
			default:
				// Keep going
			}
		}
	}()

	log.Printf("Waiting for up to %d seconds for async job %s", timeout, jobId)
	select {
	case err := <-result:
		return err
	case <-time.After(timeout):
		err := fmt.Errorf("Timeout while waiting to for async job to finish")
		return err
	}
}

// waitForAsyncJob simply blocks until the virtual machine is in the
// specified state.
func (c CloudStackClient) WaitForVirtualMachineState(vmid string, wanted_state string, timeout time.Duration) error {
	done := make(chan struct{})
	defer close(done)

	result := make(chan error, 1)
	go func() {
		attempts := 0
		for {
			attempts += 1

			log.Printf("Checking virtual machine state... (attempt: %d)", attempts)
			_, current_state, err := c.VirtualMachineState(vmid)
			if err != nil {
				result <- err
				return
			}

			// check what the real state will be.
			if current_state == wanted_state {
				result <- nil
				return
			}

			// Wait 3 seconds in between
			time.Sleep(3 * time.Second)

			// Verify we shouldn't exit
			select {
			case <-done:
				// We finished, so just exit the goroutine
				return
			default:
				// Keep going
			}
		}
	}()

	log.Printf("Waiting for up to %d seconds for Virtual Machine state to converge", timeout)
	select {
	case err := <-result:
		return err
	case <-time.After(timeout):
		err := fmt.Errorf("Timeout while waiting to for Virtual Machine to converge")
		return err
	}
}
<commit_msg>Add some debug printout<commit_after>package gopherstack

import (
	"fmt"
	"log"
	"time"
)

// waitForAsyncJob simply blocks until the the asynchronous job has
// executed or has timed out.
func (c CloudStackClient) WaitForAsyncJob(jobId string, timeout time.Duration) error {
	done := make(chan struct{})
	defer close(done)

	result := make(chan error, 1)
	go func() {
		attempts := 0
		for {
			attempts += 1

			log.Printf("Checking async job status... (attempt: %d)", attempts)
			status, err := c.QueryAsyncJobResult(jobId)
			if err != nil {
				result <- err
				return
			}

			// job is completed, we can exit
			if status == 1 {
				result <- nil
				return
			}

			// Wait 3 seconds in between
			time.Sleep(3 * time.Second)

			// Verify we shouldn't exit
			select {
			case <-done:
				// We finished, so just exit the goroutine
				return
			default:
				// Keep going
			}
		}
	}()

	log.Printf("Waiting for up to %d seconds for async job %s", timeout, jobId)
	select {
	case err := <-result:
		return err
	case <-time.After(timeout):
		err := fmt.Errorf("Timeout while waiting to for async job to finish")
		return err
	}
}

// waitForAsyncJob simply blocks until the virtual machine is in the
// specified state.
func (c CloudStackClient) WaitForVirtualMachineState(vmid string, wanted_state string, timeout time.Duration) error {
	done := make(chan struct{})
	defer close(done)

	result := make(chan error, 1)
	go func() {
		attempts := 0
		for {
			attempts += 1

			log.Printf("Checking virtual machine state... (attempt: %d)", attempts)
			_, current_state, err := c.VirtualMachineState(vmid)
			if err != nil {
				result <- err
				return
			}

			// check what the real state will be.
			log.Printf("current_state: %s", current_state)
			log.Printf("wanted_state:  %s", wanted_state)
			if current_state == wanted_state {
				result <- nil
				return
			}

			// Wait 3 seconds in between
			time.Sleep(3 * time.Second)

			// Verify we shouldn't exit
			select {
			case <-done:
				// We finished, so just exit the goroutine
				return
			default:
				// Keep going
			}
		}
	}()

	log.Printf("Waiting for up to %d seconds for Virtual Machine state to converge", timeout)
	select {
	case err := <-result:
		return err
	case <-time.After(timeout):
		err := fmt.Errorf("Timeout while waiting to for Virtual Machine to converge")
		return err
	}
}
<|endoftext|>


<commit_before>package e2e_test

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/require"

	e2e "github.com/tendermint/tendermint/test/e2e/pkg"
	"github.com/tendermint/tendermint/types"
)

// Tests that validator sets are available and correct according to
// scheduled validator updates.
func TestValidator_Sets(t *testing.T) {
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode == e2e.ModeSeed {
			return
		}

		client, err := node.Client()
		require.NoError(t, err)
		status, err := client.Status(ctx)
		require.NoError(t, err)

		first := status.SyncInfo.EarliestBlockHeight
		last := status.SyncInfo.LatestBlockHeight

		// skip first block if node is pruning blocks, to avoid race conditions
		if node.RetainBlocks > 0 {
			first++
		}

		valSchedule := newValidatorSchedule(*node.Testnet)
		valSchedule.Increment(first - node.Testnet.InitialHeight)

		for h := first; h <= last; h++ {
			validators := []*types.Validator{}
			perPage := 100
			for page := 1; ; page++ {
				resp, err := client.Validators(ctx, &(h), &(page), &perPage)
				require.NoError(t, err)
				validators = append(validators, resp.Validators...)
				if len(validators) == resp.Total {
					break
				}
			}
			require.Equal(t, valSchedule.Set.Validators, validators,
				"incorrect validator set at height %v", h)
			valSchedule.Increment(1)
		}
	})
}

// Tests that a validator proposes blocks when it's supposed to. It tolerates some
// missed blocks, e.g. due to testnet perturbations.
func TestValidator_Propose(t *testing.T) {
	blocks := fetchBlockChain(t)
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode != e2e.ModeValidator {
			return
		}
		address := node.PrivvalKey.PubKey().Address()
		valSchedule := newValidatorSchedule(*node.Testnet)

		expectCount := 0
		proposeCount := 0
		for _, block := range blocks {
			if bytes.Equal(valSchedule.Set.Proposer.Address, address) {
				expectCount++
				if bytes.Equal(block.ProposerAddress, address) {
					proposeCount++
				}
			}
			valSchedule.Increment(1)
		}

		require.False(t, proposeCount == 0 && expectCount > 0,
			"node did not propose any blocks (expected %v)", expectCount)
		if expectCount > 3 {
			require.GreaterOrEqual(t, proposeCount, 3, "validator didn't propose even 3 blocks")
		} else {
			require.Equal(t, proposeCount, expectCount, "validator missed proposing blocks (proposed %v, expected %v)",
				proposeCount, expectCount)
		}
	})
}

// Tests that a validator signs blocks when it's supposed to. It tolerates some
// missed blocks, e.g. due to testnet perturbations.
func TestValidator_Sign(t *testing.T) {
	blocks := fetchBlockChain(t)
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode != e2e.ModeValidator {
			return
		}
		address := node.PrivvalKey.PubKey().Address()
		valSchedule := newValidatorSchedule(*node.Testnet)

		expectCount := 0
		signCount := 0
		for _, block := range blocks[1:] { // Skip first block, since it has no signatures
			signed := false
			for _, sig := range block.LastCommit.Signatures {
				if bytes.Equal(sig.ValidatorAddress, address) {
					signed = true
					break
				}
			}
			if valSchedule.Set.HasAddress(address) {
				expectCount++
				if signed {
					signCount++
				}
			} else {
				require.False(t, signed, "unexpected signature for block %v", block.LastCommit.Height)
			}
			valSchedule.Increment(1)
		}

		require.False(t, signCount == 0 && expectCount > 0,
			"validator did not sign any blocks (expected %v)", expectCount)
		if expectCount > 3 {
			require.GreaterOrEqual(t, signCount, 3, "validator didn't sign even 3 blocks")
		} else {
			require.Equal(t, signCount, expectCount, "validator missed signing blocks (signed %v, expected %v)",
				signCount, expectCount)
		}
	})
}

// validatorSchedule is a validator set iterator, which takes into account
// validator set updates.
type validatorSchedule struct {
	Set     *types.ValidatorSet
	height  int64
	updates map[int64]map[*e2e.Node]int64
}

func newValidatorSchedule(testnet e2e.Testnet) *validatorSchedule {
	valMap := testnet.Validators                  // genesis validators
	if v, ok := testnet.ValidatorUpdates[0]; ok { // InitChain validators
		valMap = v
	}
	return &validatorSchedule{
		height:  testnet.InitialHeight,
		Set:     types.NewValidatorSet(makeVals(valMap)),
		updates: testnet.ValidatorUpdates,
	}
}

func (s *validatorSchedule) Increment(heights int64) {
	for i := int64(0); i < heights; i++ {
		s.height++
		if s.height > 2 {
			// validator set updates are offset by 2, since they only take effect
			// two blocks after they're returned.
			if update, ok := s.updates[s.height-2]; ok {
				if err := s.Set.UpdateWithChangeSet(makeVals(update)); err != nil {
					panic(err)
				}
			}
		}
		s.Set.IncrementProposerPriority(1)
	}
}

func makeVals(valMap map[*e2e.Node]int64) []*types.Validator {
	vals := make([]*types.Validator, 0, len(valMap))
	for node, power := range valMap {
		vals = append(vals, types.NewValidator(node.PrivvalKey.PubKey(), power))
	}
	return vals
}
<commit_msg>test/e2e: increase validator tolerances (#6037)<commit_after>package e2e_test

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/require"

	e2e "github.com/tendermint/tendermint/test/e2e/pkg"
	"github.com/tendermint/tendermint/types"
)

// Tests that validator sets are available and correct according to
// scheduled validator updates.
func TestValidator_Sets(t *testing.T) {
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode == e2e.ModeSeed {
			return
		}

		client, err := node.Client()
		require.NoError(t, err)
		status, err := client.Status(ctx)
		require.NoError(t, err)

		first := status.SyncInfo.EarliestBlockHeight
		last := status.SyncInfo.LatestBlockHeight

		// skip first block if node is pruning blocks, to avoid race conditions
		if node.RetainBlocks > 0 {
			first++
		}

		valSchedule := newValidatorSchedule(*node.Testnet)
		valSchedule.Increment(first - node.Testnet.InitialHeight)

		for h := first; h <= last; h++ {
			validators := []*types.Validator{}
			perPage := 100
			for page := 1; ; page++ {
				resp, err := client.Validators(ctx, &(h), &(page), &perPage)
				require.NoError(t, err)
				validators = append(validators, resp.Validators...)
				if len(validators) == resp.Total {
					break
				}
			}
			require.Equal(t, valSchedule.Set.Validators, validators,
				"incorrect validator set at height %v", h)
			valSchedule.Increment(1)
		}
	})
}

// Tests that a validator proposes blocks when it's supposed to. It tolerates some
// missed blocks, e.g. due to testnet perturbations.
func TestValidator_Propose(t *testing.T) {
	blocks := fetchBlockChain(t)
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode != e2e.ModeValidator {
			return
		}
		address := node.PrivvalKey.PubKey().Address()
		valSchedule := newValidatorSchedule(*node.Testnet)

		expectCount := 0
		proposeCount := 0
		for _, block := range blocks {
			if bytes.Equal(valSchedule.Set.Proposer.Address, address) {
				expectCount++
				if bytes.Equal(block.ProposerAddress, address) {
					proposeCount++
				}
			}
			valSchedule.Increment(1)
		}

		require.False(t, proposeCount == 0 && expectCount > 0,
			"node did not propose any blocks (expected %v)", expectCount)
		if expectCount > 5 {
			require.GreaterOrEqual(t, proposeCount, 3, "validator didn't propose even 3 blocks")
		}
	})
}

// Tests that a validator signs blocks when it's supposed to. It tolerates some
// missed blocks, e.g. due to testnet perturbations.
func TestValidator_Sign(t *testing.T) {
	blocks := fetchBlockChain(t)
	testNode(t, func(t *testing.T, node e2e.Node) {
		if node.Mode != e2e.ModeValidator {
			return
		}
		address := node.PrivvalKey.PubKey().Address()
		valSchedule := newValidatorSchedule(*node.Testnet)

		expectCount := 0
		signCount := 0
		for _, block := range blocks[1:] { // Skip first block, since it has no signatures
			signed := false
			for _, sig := range block.LastCommit.Signatures {
				if bytes.Equal(sig.ValidatorAddress, address) {
					signed = true
					break
				}
			}
			if valSchedule.Set.HasAddress(address) {
				expectCount++
				if signed {
					signCount++
				}
			} else {
				require.False(t, signed, "unexpected signature for block %v", block.LastCommit.Height)
			}
			valSchedule.Increment(1)
		}

		require.False(t, signCount == 0 && expectCount > 0,
			"validator did not sign any blocks (expected %v)", expectCount)
		if expectCount > 7 {
			require.GreaterOrEqual(t, signCount, 3, "validator didn't sign even 3 blocks (expected %v)", expectCount)
		}
	})
}

// validatorSchedule is a validator set iterator, which takes into account
// validator set updates.
type validatorSchedule struct {
	Set     *types.ValidatorSet
	height  int64
	updates map[int64]map[*e2e.Node]int64
}

func newValidatorSchedule(testnet e2e.Testnet) *validatorSchedule {
	valMap := testnet.Validators                  // genesis validators
	if v, ok := testnet.ValidatorUpdates[0]; ok { // InitChain validators
		valMap = v
	}
	return &validatorSchedule{
		height:  testnet.InitialHeight,
		Set:     types.NewValidatorSet(makeVals(valMap)),
		updates: testnet.ValidatorUpdates,
	}
}

func (s *validatorSchedule) Increment(heights int64) {
	for i := int64(0); i < heights; i++ {
		s.height++
		if s.height > 2 {
			// validator set updates are offset by 2, since they only take effect
			// two blocks after they're returned.
			if update, ok := s.updates[s.height-2]; ok {
				if err := s.Set.UpdateWithChangeSet(makeVals(update)); err != nil {
					panic(err)
				}
			}
		}
		s.Set.IncrementProposerPriority(1)
	}
}

func makeVals(valMap map[*e2e.Node]int64) []*types.Validator {
	vals := make([]*types.Validator, 0, len(valMap))
	for node, power := range valMap {
		vals = append(vals, types.NewValidator(node.PrivvalKey.PubKey(), power))
	}
	return vals
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go 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 main

import (
	"bufio";
	"flag";
	"fmt";
	"http";
	"io";
	"net";
	"os";
)

// hello world, the web server
func HelloServer(c *http.Conn, req *http.Request) {
	io.WriteString(c, "hello, world!\n");
}

// simple counter server
type Counter struct {
	n int;
}

func (ctr *Counter) ServeHTTP(c *http.Conn, req *http.Request) {
	fmt.Fprintf(c, "counter = %d\n", ctr.n);
	ctr.n++;
}

// simple file server
var webroot = flag.String("root", "/home/rsc", "web root directory")
func FileServer(c *http.Conn, req *http.Request) {
	c.SetHeader("content-type", "text/plain; charset=utf-8");
	path := *webroot + req.Url.Path;	// TODO: insecure: use os.CleanName
	fd, err := os.Open(path, os.O_RDONLY, 0);
	if err != nil {
		c.WriteHeader(http.StatusNotFound);
		fmt.Fprintf(c, "open %s: %v\n", path, err);
		return;
	}
	n, err1 := io.Copy(fd, c);
	fmt.Fprintf(c, "[%d bytes]\n", n);
}

// simple flag server
var booleanflag = flag.Bool("boolean", true, "another flag for testing")
func FlagServer(c *http.Conn, req *http.Request) {
	c.SetHeader("content-type", "text/plain; charset=utf-8");
	fmt.Fprint(c, "Flags:\n");
	flag.VisitAll(func (f *flag.Flag) {
		if f.Value.String() != f.DefValue {
			fmt.Fprintf(c, "%s = %s [default = %s]\n", f.Name, f.Value.String(), f.DefValue);
		} else {
			fmt.Fprintf(c, "%s = %s\n", f.Name, f.Value.String());
		}
	});
}

// simple argument server
func ArgServer(c *http.Conn, req *http.Request) {
	for i, s := range sys.Args {
		fmt.Fprint(c, s, " ");
	}
}

// a channel (just for the fun of it)
type Chan chan int

func ChanCreate() Chan {
	c := make(Chan);
	go func(c Chan) {
		for x := 0;; x++ {
			c <- x
		}
	}(c);
	return c;
}

func (ch Chan) ServeHTTP(c *http.Conn, req *http.Request) {
	io.WriteString(c, fmt.Sprintf("channel send #%d\n", <-ch));
}

func main() {
	flag.Parse();
	http.Handle("/counter", new(Counter));
	http.Handle("/go/", http.HandlerFunc(FileServer));
	http.Handle("/flags/", http.HandlerFunc(FlagServer));
	http.Handle("/args/", http.HandlerFunc(ArgServer));
	http.Handle("/go/hello", http.HandlerFunc(HelloServer));
	http.Handle("/chan", ChanCreate());
	err := http.ListenAndServe(":12345", nil);
	if err != nil {
		panic("ListenAndServe: ", err.String())
	}
}

<commit_msg>drop trailing slashes - missed comment from last review<commit_after>// Copyright 2009 The Go 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 main

import (
	"bufio";
	"flag";
	"fmt";
	"http";
	"io";
	"net";
	"os";
)

// hello world, the web server
func HelloServer(c *http.Conn, req *http.Request) {
	io.WriteString(c, "hello, world!\n");
}

// simple counter server
type Counter struct {
	n int;
}

func (ctr *Counter) ServeHTTP(c *http.Conn, req *http.Request) {
	fmt.Fprintf(c, "counter = %d\n", ctr.n);
	ctr.n++;
}

// simple file server
var webroot = flag.String("root", "/home/rsc", "web root directory")
func FileServer(c *http.Conn, req *http.Request) {
	c.SetHeader("content-type", "text/plain; charset=utf-8");
	path := *webroot + req.Url.Path;	// TODO: insecure: use os.CleanName
	fd, err := os.Open(path, os.O_RDONLY, 0);
	if err != nil {
		c.WriteHeader(http.StatusNotFound);
		fmt.Fprintf(c, "open %s: %v\n", path, err);
		return;
	}
	n, err1 := io.Copy(fd, c);
	fmt.Fprintf(c, "[%d bytes]\n", n);
}

// simple flag server
var booleanflag = flag.Bool("boolean", true, "another flag for testing")
func FlagServer(c *http.Conn, req *http.Request) {
	c.SetHeader("content-type", "text/plain; charset=utf-8");
	fmt.Fprint(c, "Flags:\n");
	flag.VisitAll(func (f *flag.Flag) {
		if f.Value.String() != f.DefValue {
			fmt.Fprintf(c, "%s = %s [default = %s]\n", f.Name, f.Value.String(), f.DefValue);
		} else {
			fmt.Fprintf(c, "%s = %s\n", f.Name, f.Value.String());
		}
	});
}

// simple argument server
func ArgServer(c *http.Conn, req *http.Request) {
	for i, s := range sys.Args {
		fmt.Fprint(c, s, " ");
	}
}

// a channel (just for the fun of it)
type Chan chan int

func ChanCreate() Chan {
	c := make(Chan);
	go func(c Chan) {
		for x := 0;; x++ {
			c <- x
		}
	}(c);
	return c;
}

func (ch Chan) ServeHTTP(c *http.Conn, req *http.Request) {
	io.WriteString(c, fmt.Sprintf("channel send #%d\n", <-ch));
}

func main() {
	flag.Parse();
	http.Handle("/counter", new(Counter));
	http.Handle("/go/", http.HandlerFunc(FileServer));
	http.Handle("/flags", http.HandlerFunc(FlagServer));
	http.Handle("/args", http.HandlerFunc(ArgServer));
	http.Handle("/go/hello", http.HandlerFunc(HelloServer));
	http.Handle("/chan", ChanCreate());
	err := http.ListenAndServe(":12345", nil);
	if err != nil {
		panic("ListenAndServe: ", err.String())
	}
}

<|endoftext|>


<commit_before><commit_msg>fix: map<commit_after><|endoftext|>


<commit_before>package main

import "fmt"
import "time"

func main() {

	numOfSundays := 0

	for y := 1901; y <= 2000; y++ {
		for m := 1; m <= 12; m++ {
			t := time.Date(y, time.Month(m), 1, 1, 0, 0, 0, time.UTC)
			if t.Weekday().String() == "Sunday" {
				numOfSundays++
			}
		}
	}

	fmt.Println(numOfSundays)
}
<commit_msg>add concurrent calculation<commit_after>package main

import "fmt"
import "time"

func calculateInYear(year int, sundays chan int) {
	sum := 0
	for m := 1; m <= 12; m++ {
		t := time.Date(year, time.Month(m), 1, 0, 0, 0, 0, time.UTC)
		if t.Weekday().String() == "Sunday" {
			sum++
		}
	}
	sundays <- sum
}

func main() {

	numOfSundays := 0
	sundaysChan := make(chan int)

	for y := 1901; y <= 2000; y++ {
		go calculateInYear(y, sundaysChan)
	}

	for y := 1901; y <= 2000; y++ {
		numOfSundays = numOfSundays + <-sundaysChan
	}

	fmt.Println(numOfSundays)
}
<|endoftext|>


<commit_before>package main

import (
	"html/template"
	"io/ioutil"
	"net/http"
)

const (
	viewPath = "/view/"
	editPath = "/edit/"
	savePath = "/save/"
)

type Page struct {
	Title string
	Body  []byte
}

func (p *Page) save() error {
	filename := p.Title + ".txt"
	return ioutil.WriteFile(filename, p.Body, 0600)
}

func loadPage(title string) (*Page, error) {
	filename := title + ".txt"
	body, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	return &Page{title, body}, err
}

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	t, err := template.ParseFiles(tmpl)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	err = t.Execute(w, p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func viewHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(viewPath):]
	page, err := loadPage(title)
	if err != nil {
		http.Redirect(w, r, editPath+title, http.StatusFound)
		return
	}

	renderTemplate(w, "view.html", page)
}

func editHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(editPath):]
	page, err := loadPage(title)
	if err != nil {
		page = &Page{title, nil}
	}

	renderTemplate(w, "edit.html", page)
}

func saveHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(savePath):]
	body := r.FormValue("body")

	p := &Page{title, []byte(body)}
	p.save()

	http.Redirect(w, r, "/view/"+title, http.StatusFound)
}

func main() {
	page := &Page{"TestPage", []byte("This is a sample page.")}
	page.save()

	http.HandleFunc(viewPath, viewHandler)
	http.HandleFunc(editPath, editHandler)
	http.HandleFunc(savePath, saveHandler)
	http.ListenAndServe("localhost:8080", nil)
}
<commit_msg>Add error checking to saveHandler<commit_after>package main

import (
	"html/template"
	"io/ioutil"
	"net/http"
)

const (
	viewPath = "/view/"
	editPath = "/edit/"
	savePath = "/save/"
)

type Page struct {
	Title string
	Body  []byte
}

func (p *Page) save() error {
	filename := p.Title + ".txt"
	return ioutil.WriteFile(filename, p.Body, 0600)
}

func loadPage(title string) (*Page, error) {
	filename := title + ".txt"
	body, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	return &Page{title, body}, err
}

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	t, err := template.ParseFiles(tmpl)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	err = t.Execute(w, p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func viewHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(viewPath):]
	page, err := loadPage(title)
	if err != nil {
		http.Redirect(w, r, editPath+title, http.StatusFound)
		return
	}

	renderTemplate(w, "view.html", page)
}

func editHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(editPath):]
	page, err := loadPage(title)
	if err != nil {
		page = &Page{title, nil}
	}

	renderTemplate(w, "edit.html", page)
}

func saveHandler(w http.ResponseWriter, r *http.Request) {
	title := r.URL.Path[len(savePath):]
	body := r.FormValue("body")

	p := &Page{title, []byte(body)}
	err := p.save()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	http.Redirect(w, r, "/view/"+title, http.StatusFound)
}

func main() {
	page := &Page{"TestPage", []byte("This is a sample page.")}
	page.save()

	http.HandleFunc(viewPath, viewHandler)
	http.HandleFunc(editPath, editHandler)
	http.HandleFunc(savePath, saveHandler)
	http.ListenAndServe("localhost:8080", nil)
}
<|endoftext|>


<commit_before>package main

import (
	// "errors"
	"flag"
	"fmt"
	"gopkg.in/yaml.v2"
	"html/template"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"regexp"
	// "strings"
)

var (
	addr       = flag.Bool("addr", false, "find open address and print to final-port.txt")
	configfile = flag.String("configfile", "config.yaml", "path and filename of the config file")
)

type Config struct {
	// First letter of variables need to be capital letter
	Template_directory string
	Data_directory     string
}

var config Config

// config.Data_directory
var TMPL_DIR = "./templates/"
var DATA_DIR = "./data/"

type Page struct {
	Title string
	Body  []byte
}

var templates *template.Template

var validPath = regexp.MustCompile("^/(edit|save|view|list)/([a-zA-Z0-9]+)$")

func (p *Page) save() error {
	filename := DATA_DIR + p.Title + ".txt"
	return ioutil.WriteFile(filename, p.Body, 0600)
}

// func getTitle(w http.ResponseWriter, r *http.Request) (string, error) {
// 	m := validPath.FindStringSubmatch(r.URL.Path)
// 	if m == nil {
// 		http.NotFound(w, r)
// 		return "", errors.New("Invalid Page Title")
// 	}
// 	return m[2], nil
// }

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	err := templates.ExecuteTemplate(w, tmpl+".html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
// 	t, err := template.ParseFiles(tmpl + ".html")
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 		return
// 	}
// 	err = t.Execute(w, p)
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 	}
// }

func loadPage(title string) (*Page, error) {
	filename := DATA_DIR + title + ".txt"
	body, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return &Page{Title: title, Body: body}, nil
}

func editHandler(w http.ResponseWriter, r *http.Request, title string) {
	p, err := loadPage(title)
	if err != nil {
		p = &Page{Title: title}
	}
	renderTemplate(w, "edit", p)
}

func viewHandler(w http.ResponseWriter, r *http.Request, title string) {
	p, err := loadPage(title)
	if err != nil {
		http.Redirect(w, r, "/edit/"+title, http.StatusFound)
		return
	}
	renderTemplate(w, "view", p)
}

func saveHandler(w http.ResponseWriter, r *http.Request, title string) {
	body := r.FormValue("body")
	p := &Page{Title: title, Body: []byte(body)}
	err := p.save()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	http.Redirect(w, r, "/view/"+title, http.StatusFound)
}

func listHandler(w http.ResponseWriter, r *http.Request, title string) {
	datafiles, err := ioutil.ReadDir(DATA_DIR)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	// layoutData := struct {
	// 	datafiles []string
	// }
	for _, f := range datafiles {
		fmt.Println(f.Name())

	}

	// renderTemplate(w, "list", datafiles)
	err = templates.ExecuteTemplate(w, "list.html", datafiles)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Similar to Decorator in Python
func makeHandler(fn func(http.ResponseWriter, *http.Request, string)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		m := validPath.FindStringSubmatch(r.URL.Path)
		if m == nil {
			http.NotFound(w, r)
			return
		}
		fn(w, r, m[2])
	}
}

func main() {
	flag.Parse()
	// prepare handler
	http.HandleFunc("/list/", makeHandler(listHandler))
	http.HandleFunc("/view/", makeHandler(viewHandler))
	http.HandleFunc("/edit/", makeHandler(editHandler))
	http.HandleFunc("/save/", makeHandler(saveHandler))

	config = Config{}

	// Load configfile and configure template
	if len(*configfile) > 0 {
		fmt.Println("config file => " + *configfile)
		source, err := ioutil.ReadFile(*configfile)
		fmt.Println(string(source))
		if err != nil {
			panic(err)
		}
		// decode the yaml source
		err = yaml.Unmarshal(source, &config)
		if err != nil {
			panic(err)
		}
		TMPL_DIR = config.Template_directory
		DATA_DIR = config.Data_directory
	}

	templates = template.Must(template.ParseFiles(TMPL_DIR+"edit.html", TMPL_DIR+"view.html", TMPL_DIR+"list.html"))

	if *addr {
		l, err := net.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			log.Fatal(err)
		}
		err = ioutil.WriteFile("final-port.txt", []byte(l.Addr().String()), 0644)
		if err != nil {
			log.Fatal(err)
		}
		s := &http.Server{}
		s.Serve(l)
		return
	}
	http.ListenAndServe(":8080", nil)
}
<commit_msg>create basic middleware using Alice<commit_after>package main

import (
	// "errors"
	// "strings"
	"flag"
	"fmt"
	"github.com/justinas/alice"
	"gopkg.in/yaml.v2"
	"html/template"
	// "io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"regexp"
	"time"
)

var (
	addr       = flag.Bool("addr", false, "find open address and print to final-port.txt")
	configfile = flag.String("configfile", "config.yaml", "path and filename of the config file")
)

type Config struct {
	// First letter of variables need to be capital letter
	Template_directory string
	Data_directory     string
}

var config Config

// config.Data_directory
var TMPL_DIR = "./templates/"
var DATA_DIR = "./data/"

type Page struct {
	Title string
	Body  []byte
}

var templates *template.Template

var validPath = regexp.MustCompile("^/(edit|save|view|list)/([a-zA-Z0-9]*)$")

func (p *Page) save() error {
	filename := DATA_DIR + p.Title + ".txt"
	return ioutil.WriteFile(filename, p.Body, 0600)
}

// func getTitle(w http.ResponseWriter, r *http.Request) (string, error) {
// 	m := validPath.FindStringSubmatch(r.URL.Path)
// 	if m == nil {
// 		http.NotFound(w, r)
// 		return "", errors.New("Invalid Page Title")
// 	}
// 	return m[2], nil
// }

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	err := templates.ExecuteTemplate(w, tmpl+".html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
// 	t, err := template.ParseFiles(tmpl + ".html")
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 		return
// 	}
// 	err = t.Execute(w, p)
// 	if err != nil {
// 		http.Error(w, err.Error(), http.StatusInternalServerError)
// 	}
// }

func loadPage(title string) (*Page, error) {
	filename := DATA_DIR + title + ".txt"
	body, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return &Page{Title: title, Body: body}, nil
}

func editHandler(w http.ResponseWriter, r *http.Request, title string) {
	p, err := loadPage(title)
	if err != nil {
		p = &Page{Title: title}
	}
	renderTemplate(w, "edit", p)
}

func viewHandler(w http.ResponseWriter, r *http.Request, title string) {
	p, err := loadPage(title)
	if err != nil {
		http.Redirect(w, r, "/edit/"+title, http.StatusFound)
		return
	}
	renderTemplate(w, "view", p)
}

func saveHandler(w http.ResponseWriter, r *http.Request, title string) {
	body := r.FormValue("body")
	p := &Page{Title: title, Body: []byte(body)}
	err := p.save()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	http.Redirect(w, r, "/view/"+title, http.StatusFound)
}

func listHandler(w http.ResponseWriter, r *http.Request, title string) {
	datafiles, err := ioutil.ReadDir(DATA_DIR)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	// layoutData := struct {
	// 	datafiles []string
	// }
	for _, f := range datafiles {
		fmt.Println(f.Name())
	}

	// renderTemplate(w, "list", datafiles)
	err = templates.ExecuteTemplate(w, "list.html", datafiles)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Similar to Decorator in Python
func makeHandler(fn func(http.ResponseWriter, *http.Request, string)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		m := validPath.FindStringSubmatch(r.URL.Path)
		if m == nil {
			http.NotFound(w, r)
			return
		}
		fn(w, r, m[2])
	}
}

func loggingHandler(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		t1 := time.Now()
		next.ServeHTTP(w, r)
		t2 := time.Now()
		log.Printf("[%s] %q %v\n", r.Method, r.URL.String(), t2.Sub(t1))
	}
	return http.HandlerFunc(fn)
}

func aboutHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "You are on the about page.")
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Welcome to Weekee!")
}

func recoverHandler(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				log.Printf("panic: %+v", err)
				http.Error(w, http.StatusText(500), 500)
			}
		}()
		next.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

func main() {
	flag.Parse()

	commonHandlers := alice.New(loggingHandler, recoverHandler)
	http.Handle("/about", commonHandlers.ThenFunc(aboutHandler))
	http.Handle("/", commonHandlers.ThenFunc(indexHandler))

	// prepare handler
	http.HandleFunc("/list/", makeHandler(listHandler))
	http.HandleFunc("/view/", makeHandler(viewHandler))
	http.HandleFunc("/edit/", makeHandler(editHandler))
	http.HandleFunc("/save/", makeHandler(saveHandler))

	config = Config{}

	// Load configfile and configure template
	if len(*configfile) > 0 {
		fmt.Println("config file => " + *configfile)
		source, err := ioutil.ReadFile(*configfile)
		fmt.Println(string(source))
		if err != nil {
			panic(err)
		}
		// decode the yaml source
		err = yaml.Unmarshal(source, &config)
		if err != nil {
			panic(err)
		}
		TMPL_DIR = config.Template_directory
		DATA_DIR = config.Data_directory
	}

	templates = template.Must(template.ParseFiles(TMPL_DIR+"edit.html", TMPL_DIR+"view.html", TMPL_DIR+"list.html"))

	if *addr {
		l, err := net.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			log.Fatal(err)
		}
		err = ioutil.WriteFile("final-port.txt", []byte(l.Addr().String()), 0644)
		if err != nil {
			log.Fatal(err)
		}
		s := &http.Server{}
		s.Serve(l)
		return
	}
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>// Copyright 2015 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

// wl2k provides a client-side implementation of the B2 Forwarding Protocol
// and Winlink 2000 Message Structure for transfer of messages to and from
// a Winlink 2000 Radio Email Server (RMS) gateway.
package wl2k

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// Objects implementing the MBoxHandler interface can be used to handle inbound and outbound messages for a Session.
type MBoxHandler interface {
	InboundHandler
	OutboundHandler

	// Prepare is called before any other operation in a session.
	//
	// The returned error can be used to indicate that the mailbox is
	// not ready for a new session, the error will be forwarded to the
	// remote node.
	Prepare() error
}

// An OutboundHandler offer messages that can be delivered (a proposal) to the remote node and is notified when a message is sent or defered.
type OutboundHandler interface {
	// GetOutbound should return all pending (outbound) messages addressed to (and only to) one of the fw addresses.
	//
	// No fw address implies that the remote node could be a Winlink CMS and all oubound
	// messages can be delivered through the connected node.
	GetOutbound(fw ...Address) (out []*Message)

	// SetSent should mark the the message identified by MID as successfully sent.
	//
	// If rejected is true, it implies that the remote node has already received the message.
	SetSent(MID string, rejected bool)

	// SetDeferred should mark the outbound message identified by MID as deferred.
	//
	// SetDeferred is called when the remote want's to receive the proposed message
	// (see MID) later.
	SetDeferred(MID string)
}

// An InboundHandler handles all messages that can/is sent from the remote node.
type InboundHandler interface {
	// ProcessInbound should persist/save/process all messages received (msgs) returning an error if the operation was unsuccessful.
	//
	// The error will be delivered (if possble) to the remote to indicate that an error has occurred.
	ProcessInbound(msg ...*Message) error

	// GetInboundAnswer should return a ProposalAnwer (Accept/Reject/Defer) based on the remote's message Proposal p.
	//
	// An already successfully received message (see MID) should be rejected.
	GetInboundAnswer(p Proposal) ProposalAnswer
}

// Session represents a B2F exchange session.
//
// A session should only be used once.
type Session struct {
	mycall     string
	targetcall string
	locator    string
	motd       []string

	h             MBoxHandler
	statusUpdater StatusUpdater

	// Callback when secure login password is needed
	secureLoginHandleFunc func() (password string, err error)

	master bool

	remoteSID sid
	remoteFW  []Address // Addresses the remote requests messages on behalf of
	localFW   []Address // Addresses we request messages on behalf of

	trafficStats TrafficStats

	quitReceived bool
	quitSent     bool
	remoteNoMsgs bool // True if last remote turn had no more messages

	rd *bufio.Reader

	log  *log.Logger
	pLog *log.Logger
	ua   UserAgent
}

// Struct used to hold information that is reported during B2F handshake.
//
// Non of the fields must contain a dash (-).
//
type UserAgent struct {
	Name    string
	Version string
}

type StatusUpdater interface {
	UpdateStatus(s Status)
}

// Status holds information about ongoing transfers.
type Status struct {
	Receiving        *Proposal
	Sending          *Proposal
	BytesTransferred int
	BytesTotal       int
	When             time.Time
}

// TrafficStats holds exchange message traffic statistics.
type TrafficStats struct {
	Received []string // Received message MIDs.
	Sent     []string // Sent message MIDs.
}

var StdLogger = log.New(os.Stderr, "", log.LstdFlags)
var StdUA = UserAgent{Name: "wl2kgo", Version: "0.1a"}

// Constructs a new Session object.
//
// The Handler can be nil (but no messages will be exchanged).
func NewSession(mycall, targetcall, locator string, h MBoxHandler) *Session {
	return &Session{
		mycall:     mycall,
		localFW:    []Address{AddressFromString(mycall)},
		targetcall: targetcall,
		log:        StdLogger,
		h:          h,
		pLog:       StdLogger,
		ua:         StdUA,
		locator:    locator,
	}
}

// SetMOTD sets one or more lines to be sent before handshake.
//
// The MOTD is only sent if the local node is session master.
func (s *Session) SetMOTD(line ...string) { s.motd = line }

// IsMaster sets whether this end should initiate the handshake.
func (s *Session) IsMaster(isMaster bool) { s.master = isMaster }

// RemoteSID returns the remote's SID (if available).
func (s *Session) RemoteSID() string { return string(s.remoteSID) }

// Exchange is the main method for exchanging messages with a remote over the B2F protocol.
//
// Sends outbound messages and downloads inbound messages prepared for this session.
//
// Outbound messages should be added as proposals before calling the Exchange() method.
//
// After Exchange(), messages that was accepted and delivered successfully to the RMS is
// available through a call to Sent(). Messages downloaded successfully from the RMS is
// retrieved by calling Received().
//
// The connection is closed at the end of the exchange. If the connection is closed before
// the exchange is done, is will return io.EOF.
//
// Subsequent Exchange calls on the same session is a noop.
func (s *Session) Exchange(conn net.Conn) (stats TrafficStats, err error) {
	if s.Done() {
		return stats, nil
	}

	// The given conn should always be closed after returning from this method.
	// If an error occured, echo it to the remote.
	defer func() {
		if err == nil {
			return
		}

		// In case another go-routine closes the connection...
		localEOF := strings.Contains(err.Error(), "use of closed network connection")
		if localEOF {
			err = io.EOF
		}

		if err != io.EOF {
			conn.SetDeadline(time.Now().Add(time.Minute))
			fmt.Fprintf(conn, "*** %s\r\n", err)
			conn.Close()
		}
	}()

	// Prepare mailbox handler
	if s.h != nil {
		err = s.h.Prepare()
		if err != nil {
			return
		}
	}

	s.rd = bufio.NewReader(conn)

	err = s.handshake(conn)
	if err != nil {
		return
	}

	for myTurn := !s.master; !s.Done(); myTurn = !myTurn {
		if myTurn {
			s.quitSent, err = s.handleOutbound(conn)
		} else {
			s.quitReceived, err = s.handleInbound(conn)
		}

		if err != nil {
			return s.trafficStats, err
		}
	}

	return s.trafficStats, conn.Close()
}

// Done() returns true if either parties have existed from this session.
func (s *Session) Done() bool { return s.quitReceived || s.quitSent }

// Waits for connection to be closed, returning an error if seen on the line.
func waitRemoteHangup(conn net.Conn) error {
	conn.SetDeadline(time.Now().Add(time.Minute))

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		line := scanner.Text()

		if err := errLine(line); err != nil {
			conn.Close()
			return err
		}
		log.Println(line)
	}
	return scanner.Err()
}

func remoteErr(str string) error {
	if !strings.HasPrefix(str, "***") {
		return nil
	}

	idx := strings.LastIndex(str, "*")
	if idx+1 >= len(str) {
		return nil
	}

	return fmt.Errorf(strings.TrimSpace(str[idx+1:]))
}

// Mycall returns this stations call sign.
func (s *Session) Mycall() string { return s.mycall }

// Targetcall returns the remote stations call sign (if known).
func (s *Session) Targetcall() string { return s.targetcall }

// SetSecureLoginHandleFunc registers a callback function used to prompt for password when a secure login challenge is received.
func (s *Session) SetSecureLoginHandleFunc(f func() (password string, err error)) {
	s.secureLoginHandleFunc = f
}

// This method returns the call signs the remote is requesting traffic on behalf of. The call signs are not available until
// the handshake is done.
//
// It will typically be the call sign of the remote P2P station and empty when the remote is a Winlink CMS.
func (s *Session) RemoteForwarders() []Address { return s.remoteFW }

// AddAuxiliaryAddress adds one or more addresses to request messages on behalf of.
//
// Currently the Winlink System only support requesting messages for call signs, not full email addresses.
func (s *Session) AddAuxiliaryAddress(aux ...Address) { s.localFW = append(s.localFW, aux...) }

// Set callback for status updates on receiving / sending messages
func (s *Session) SetStatusUpdater(updater StatusUpdater) { s.statusUpdater = updater }

// Sets custom logger.
func (s *Session) SetLogger(logger *log.Logger) {
	if logger == nil {
		logger = StdLogger
	}
	s.log = logger
	s.pLog = logger

}

// Set this session's user agent
func (s *Session) SetUserAgent(ua UserAgent) { s.ua = ua }

// Get this session's user agent
func (s *Session) UserAgent() UserAgent { return s.ua }

func (s *Session) outbound() []*Proposal {
	if s.h == nil {
		return []*Proposal{}
	}

	msgs := s.h.GetOutbound(s.remoteFW...)
	props := make([]*Proposal, 0, len(msgs))

	for _, m := range msgs {
		prop, err := m.Proposal()
		if err != nil {
			// TODO: This should result in an error somewhere
			s.log.Printf("Unable to prepare proposal for '%s'. Corrupt message? Skipping...", prop.MID())
			continue
		}

		props = append(props, prop)
	}
	return props
}
<commit_msg>wl2k: Fix nil slice in stats<commit_after>// Copyright 2015 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

// wl2k provides a client-side implementation of the B2 Forwarding Protocol
// and Winlink 2000 Message Structure for transfer of messages to and from
// a Winlink 2000 Radio Email Server (RMS) gateway.
package wl2k

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// Objects implementing the MBoxHandler interface can be used to handle inbound and outbound messages for a Session.
type MBoxHandler interface {
	InboundHandler
	OutboundHandler

	// Prepare is called before any other operation in a session.
	//
	// The returned error can be used to indicate that the mailbox is
	// not ready for a new session, the error will be forwarded to the
	// remote node.
	Prepare() error
}

// An OutboundHandler offer messages that can be delivered (a proposal) to the remote node and is notified when a message is sent or defered.
type OutboundHandler interface {
	// GetOutbound should return all pending (outbound) messages addressed to (and only to) one of the fw addresses.
	//
	// No fw address implies that the remote node could be a Winlink CMS and all oubound
	// messages can be delivered through the connected node.
	GetOutbound(fw ...Address) (out []*Message)

	// SetSent should mark the the message identified by MID as successfully sent.
	//
	// If rejected is true, it implies that the remote node has already received the message.
	SetSent(MID string, rejected bool)

	// SetDeferred should mark the outbound message identified by MID as deferred.
	//
	// SetDeferred is called when the remote want's to receive the proposed message
	// (see MID) later.
	SetDeferred(MID string)
}

// An InboundHandler handles all messages that can/is sent from the remote node.
type InboundHandler interface {
	// ProcessInbound should persist/save/process all messages received (msgs) returning an error if the operation was unsuccessful.
	//
	// The error will be delivered (if possble) to the remote to indicate that an error has occurred.
	ProcessInbound(msg ...*Message) error

	// GetInboundAnswer should return a ProposalAnwer (Accept/Reject/Defer) based on the remote's message Proposal p.
	//
	// An already successfully received message (see MID) should be rejected.
	GetInboundAnswer(p Proposal) ProposalAnswer
}

// Session represents a B2F exchange session.
//
// A session should only be used once.
type Session struct {
	mycall     string
	targetcall string
	locator    string
	motd       []string

	h             MBoxHandler
	statusUpdater StatusUpdater

	// Callback when secure login password is needed
	secureLoginHandleFunc func() (password string, err error)

	master bool

	remoteSID sid
	remoteFW  []Address // Addresses the remote requests messages on behalf of
	localFW   []Address // Addresses we request messages on behalf of

	trafficStats TrafficStats

	quitReceived bool
	quitSent     bool
	remoteNoMsgs bool // True if last remote turn had no more messages

	rd *bufio.Reader

	log  *log.Logger
	pLog *log.Logger
	ua   UserAgent
}

// Struct used to hold information that is reported during B2F handshake.
//
// Non of the fields must contain a dash (-).
//
type UserAgent struct {
	Name    string
	Version string
}

type StatusUpdater interface {
	UpdateStatus(s Status)
}

// Status holds information about ongoing transfers.
type Status struct {
	Receiving        *Proposal
	Sending          *Proposal
	BytesTransferred int
	BytesTotal       int
	When             time.Time
}

// TrafficStats holds exchange message traffic statistics.
type TrafficStats struct {
	Received []string // Received message MIDs.
	Sent     []string // Sent message MIDs.
}

var StdLogger = log.New(os.Stderr, "", log.LstdFlags)
var StdUA = UserAgent{Name: "wl2kgo", Version: "0.1a"}

// Constructs a new Session object.
//
// The Handler can be nil (but no messages will be exchanged).
func NewSession(mycall, targetcall, locator string, h MBoxHandler) *Session {
	return &Session{
		mycall:     mycall,
		localFW:    []Address{AddressFromString(mycall)},
		targetcall: targetcall,
		log:        StdLogger,
		h:          h,
		pLog:       StdLogger,
		ua:         StdUA,
		locator:    locator,
		trafficStats: TrafficStats{
			Received: make([]string, 0),
			Sent:     make([]string, 0),
		},
	}
}

// SetMOTD sets one or more lines to be sent before handshake.
//
// The MOTD is only sent if the local node is session master.
func (s *Session) SetMOTD(line ...string) { s.motd = line }

// IsMaster sets whether this end should initiate the handshake.
func (s *Session) IsMaster(isMaster bool) { s.master = isMaster }

// RemoteSID returns the remote's SID (if available).
func (s *Session) RemoteSID() string { return string(s.remoteSID) }

// Exchange is the main method for exchanging messages with a remote over the B2F protocol.
//
// Sends outbound messages and downloads inbound messages prepared for this session.
//
// Outbound messages should be added as proposals before calling the Exchange() method.
//
// After Exchange(), messages that was accepted and delivered successfully to the RMS is
// available through a call to Sent(). Messages downloaded successfully from the RMS is
// retrieved by calling Received().
//
// The connection is closed at the end of the exchange. If the connection is closed before
// the exchange is done, is will return io.EOF.
//
// Subsequent Exchange calls on the same session is a noop.
func (s *Session) Exchange(conn net.Conn) (stats TrafficStats, err error) {
	if s.Done() {
		return stats, nil
	}

	// The given conn should always be closed after returning from this method.
	// If an error occured, echo it to the remote.
	defer func() {
		if err == nil {
			return
		}

		// In case another go-routine closes the connection...
		localEOF := strings.Contains(err.Error(), "use of closed network connection")
		if localEOF {
			err = io.EOF
		}

		if err != io.EOF {
			conn.SetDeadline(time.Now().Add(time.Minute))
			fmt.Fprintf(conn, "*** %s\r\n", err)
			conn.Close()
		}
	}()

	// Prepare mailbox handler
	if s.h != nil {
		err = s.h.Prepare()
		if err != nil {
			return
		}
	}

	s.rd = bufio.NewReader(conn)

	err = s.handshake(conn)
	if err != nil {
		return
	}

	for myTurn := !s.master; !s.Done(); myTurn = !myTurn {
		if myTurn {
			s.quitSent, err = s.handleOutbound(conn)
		} else {
			s.quitReceived, err = s.handleInbound(conn)
		}

		if err != nil {
			return s.trafficStats, err
		}
	}

	return s.trafficStats, conn.Close()
}

// Done() returns true if either parties have existed from this session.
func (s *Session) Done() bool { return s.quitReceived || s.quitSent }

// Waits for connection to be closed, returning an error if seen on the line.
func waitRemoteHangup(conn net.Conn) error {
	conn.SetDeadline(time.Now().Add(time.Minute))

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		line := scanner.Text()

		if err := errLine(line); err != nil {
			conn.Close()
			return err
		}
		log.Println(line)
	}
	return scanner.Err()
}

func remoteErr(str string) error {
	if !strings.HasPrefix(str, "***") {
		return nil
	}

	idx := strings.LastIndex(str, "*")
	if idx+1 >= len(str) {
		return nil
	}

	return fmt.Errorf(strings.TrimSpace(str[idx+1:]))
}

// Mycall returns this stations call sign.
func (s *Session) Mycall() string { return s.mycall }

// Targetcall returns the remote stations call sign (if known).
func (s *Session) Targetcall() string { return s.targetcall }

// SetSecureLoginHandleFunc registers a callback function used to prompt for password when a secure login challenge is received.
func (s *Session) SetSecureLoginHandleFunc(f func() (password string, err error)) {
	s.secureLoginHandleFunc = f
}

// This method returns the call signs the remote is requesting traffic on behalf of. The call signs are not available until
// the handshake is done.
//
// It will typically be the call sign of the remote P2P station and empty when the remote is a Winlink CMS.
func (s *Session) RemoteForwarders() []Address { return s.remoteFW }

// AddAuxiliaryAddress adds one or more addresses to request messages on behalf of.
//
// Currently the Winlink System only support requesting messages for call signs, not full email addresses.
func (s *Session) AddAuxiliaryAddress(aux ...Address) { s.localFW = append(s.localFW, aux...) }

// Set callback for status updates on receiving / sending messages
func (s *Session) SetStatusUpdater(updater StatusUpdater) { s.statusUpdater = updater }

// Sets custom logger.
func (s *Session) SetLogger(logger *log.Logger) {
	if logger == nil {
		logger = StdLogger
	}
	s.log = logger
	s.pLog = logger

}

// Set this session's user agent
func (s *Session) SetUserAgent(ua UserAgent) { s.ua = ua }

// Get this session's user agent
func (s *Session) UserAgent() UserAgent { return s.ua }

func (s *Session) outbound() []*Proposal {
	if s.h == nil {
		return []*Proposal{}
	}

	msgs := s.h.GetOutbound(s.remoteFW...)
	props := make([]*Proposal, 0, len(msgs))

	for _, m := range msgs {
		prop, err := m.Proposal()
		if err != nil {
			// TODO: This should result in an error somewhere
			s.log.Printf("Unable to prepare proposal for '%s'. Corrupt message? Skipping...", prop.MID())
			continue
		}

		props = append(props, prop)
	}
	return props
}
<|endoftext|>