<commit_before>package servd

import (
	"context"
	"errors"
)

// Status type of the service.
type Status int

const (
	// Created status.
	Created Status = iota

	// Running status.
	Running

	// Stopped status
	Stopped
)

// Handler is the application handler.
type Handler interface {
	Handle(context.Context) error
}

// HandleFunc is the Handler adapter.
type HandleFunc func(context.Context) error

// Handle the application run.
func (f HandleFunc) Handle(ctx context.Context) error {
	return f(ctx)
}

// Servd is the service daemon.
type Servd struct {
	Handler         Handler
	status          Status
	statSubscribers map[Status][]chan<- Status
	stop            func()
}

// Run the service.
func (s *Servd) Run() error {
	if s.Handler == nil {
		return errors.New("servd: no handler")
	}

	if s.status > Created {
		return errors.New("servd: cannot re-run")
	}

	// initialize
	s.statSubscribers = make(map[Status][]chan<- Status)

	s.changeStatus(Running)
	defer s.changeStatus(Stopped)

	ctx, cancel := context.WithCancel(context.Background())
	s.stop = cancel
	return s.handle(ctx)
}

func (s *Servd) handle(ctx context.Context) error {
	return s.Handler.Handle(ctx)
}

// Status of the service.
func (s *Servd) Status() Status {
	return s.status
}

// WaitForStatus waiting for specific status.
// It will wait until it reach or passed the expected status.
func (s *Servd) WaitForStatus(ctx context.Context, stat Status) (Status, error) {
	if s.status >= stat {
		return s.status, nil
	}

	c := make(chan Status, 1)
	s.notifyStatus(c, stat)
	defer s.cancelNotifyStatus(c, stat)

	select {
	case <-ctx.Done():
		return Created, ctx.Err()
	case <-c:
		return stat, nil
	}
}

// Stop the service.
// Return true if the service stop has been called or the app is shutting down.
func (s *Servd) Stop() bool {
	if s.stop == nil && s.status == Running {
		return false
	}

	s.stop()
	return true
}

// StopAndWait the service to be stopped.
func (s *Servd) StopAndWait(ctx context.Context) (Status, error) {
	s.Stop()
	return s.WaitForStatus(ctx, Stopped)
}

func (s *Servd) notifyStatus(c chan<- Status, stat Status) {
	s.statSubscribers[stat] = append(s.statSubscribers[stat], c)
}

func (s *Servd) cancelNotifyStatus(c chan<- Status, stat Status) {
	subs := s.statSubscribers[stat]
	i := chanIndex(subs, c)
	if i < 0 {
		return
	}

	// remove value on index i
	s.statSubscribers[stat] = append(subs[:i], subs[i+1:]...)
}

func (s *Servd) changeStatus(stat Status) {
	s.status = stat

	// notifies all subscribers of s
	subs, ok := s.statSubscribers[stat]
	if !ok {
		return
	}

	notify := func(sub chan<- Status, stat Status) {
		sub <- stat
	}
	for _, sub := range subs {
		go notify(sub, stat)
	}
}

func chanIndex(s []chan<- Status, c chan<- Status) int {
	for i, v := range s {
		if v == c {
			return i
		}
	}
	return -1
}
<commit_msg>Guard nil map<commit_after>package servd

import (
	"context"
	"errors"
)

// Status type of the service.
type Status int

const (
	// Created status.
	Created Status = iota

	// Running status.
	Running

	// Stopped status
	Stopped
)

// Handler is the application handler.
type Handler interface {
	Handle(context.Context) error
}

// HandleFunc is the Handler adapter.
type HandleFunc func(context.Context) error

// Handle the application run.
func (f HandleFunc) Handle(ctx context.Context) error {
	return f(ctx)
}

// Servd is the service daemon.
type Servd struct {
	Handler         Handler
	status          Status
	statSubscribers map[Status][]chan<- Status
	stop            func()
}

// Run the service.
func (s *Servd) Run() error {
	if s.Handler == nil {
		return errors.New("servd: no handler")
	}

	if s.status > Created {
		return errors.New("servd: cannot re-run")
	}

	s.changeStatus(Running)
	defer s.changeStatus(Stopped)

	ctx, cancel := context.WithCancel(context.Background())
	s.stop = cancel
	return s.handle(ctx)
}

func (s *Servd) handle(ctx context.Context) error {
	return s.Handler.Handle(ctx)
}

// Status of the service.
func (s *Servd) Status() Status {
	return s.status
}

// WaitForStatus waiting for specific status.
// It will wait until it reach or passed the expected status.
func (s *Servd) WaitForStatus(ctx context.Context, stat Status) (Status, error) {
	if s.status >= stat {
		return s.status, nil
	}

	c := make(chan Status, 1)
	s.notifyStatus(c, stat)
	defer s.cancelNotifyStatus(c, stat)

	select {
	case <-ctx.Done():
		return Created, ctx.Err()
	case <-c:
		return stat, nil
	}
}

// Stop the service.
// Return true if the service stop has been called or the app is shutting down.
func (s *Servd) Stop() bool {
	if s.stop == nil && s.status == Running {
		return false
	}

	s.stop()
	return true
}

// StopAndWait the service to be stopped.
func (s *Servd) StopAndWait(ctx context.Context) (Status, error) {
	s.Stop()
	return s.WaitForStatus(ctx, Stopped)
}

func (s *Servd) notifyStatus(c chan<- Status, stat Status) {
	if s.statSubscribers == nil {
		s.statSubscribers = make(map[Status][]chan<- Status)
	}

	s.statSubscribers[stat] = append(s.statSubscribers[stat], c)
}

func (s *Servd) cancelNotifyStatus(c chan<- Status, stat Status) {
	if s.statSubscribers == nil {
		return
	}

	subs := s.statSubscribers[stat]
	i := chanIndex(subs, c)
	if i < 0 {
		return
	}

	// remove value on index i
	s.statSubscribers[stat] = append(subs[:i], subs[i+1:]...)
}

func (s *Servd) changeStatus(stat Status) {
	s.status = stat

	if s.statSubscribers == nil {
		return
	}

	// notifies all subscribers of s
	subs, ok := s.statSubscribers[stat]
	if !ok {
		return
	}

	notify := func(sub chan<- Status, stat Status) {
		sub <- stat
	}
	for _, sub := range subs {
		go notify(sub, stat)
	}
}

func chanIndex(s []chan<- Status, c chan<- Status) int {
	for i, v := range s {
		if v == c {
			return i
		}
	}
	return -1
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

func usage() {
	fmt.Fprintf(os.Stderr, "\nusage: go run serve.go [optional flags]\n")
	flag.PrintDefaults()
	fmt.Println()

	os.Exit(2)
}

type myHandler struct {
	certMap map[string]tls.Certificate
}

type myCert struct {
	cert      *tls.Certificate
	touchedAt time.Time
}

func (m *myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Print debug info
	fmt.Println(r.Host)
	fmt.Println(r.Method)
	fmt.Println(r.RequestURI)
	fmt.Println(r.URL) // has many keys, such as Query
	for k, v := range r.Header {
		fmt.Println(k, v)
	}

	fmt.Println(r.Body)
	fmt.Println()
	fmt.Println()

	// End the request
	// TODO serve from hosting directory
	fmt.Fprintf(w, "Hi there, %s %q? Wow!\n\nWith Love,\n\t%s", r.Method, r.URL.Path[1:], r.Host)
}

func main() {
	flag.Usage = usage

	port := flag.Uint("port", 443, "https port")
	certsPath := flag.String("letsencrypt-path", "/etc/letsencrypt/live", "path at which an 'xyz.example.com' containing 'fullchain.pem' and 'privkey.pem' can be found")
	defaultHost := flag.String("default-hostname", "localhost.daplie.com", "the default folder to find certificates to use when no matches are found")

	flag.Parse()

	host := strings.ToLower(*defaultHost)
	// See https://groups.google.com/a/letsencrypt.org/forum/#!topic/ca-dev/l1Dd6jzWeu8
	/*
		if strings.HasPrefix("www.", host) {
			fmt.Println("TODO: 'www.' prefixed certs should be obtained for every 'example.com' domain.")
		}
		host = strings.TrimPrefix("www.", host)
	*/

	fmt.Printf("Loading Certificates %s/%s/{privkey.pem,fullchain.pem}\n", *certsPath, *defaultHost)
	privkeyPath := filepath.Join(*certsPath, *defaultHost, "privkey.pem")
	certPath := filepath.Join(*certsPath, *defaultHost, "fullchain.pem")
	defaultCert, err := tls.LoadX509KeyPair(certPath, privkeyPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't load default certificates: %s\n", err)
		os.Exit(1)
	}

	addr := ":" + strconv.Itoa(int(*port))

	conn, err := net.Listen("tcp", addr)
	if nil != err {
		fmt.Fprintf(os.Stderr, "Couldn't bind to TCP socket %q: %s\n", addr, err)
		os.Exit(1)
	}

	certMap := make(map[string]myCert)
	tlsConfig := new(tls.Config)
	tlsConfig.Certificates = []tls.Certificate{defaultCert}
	tlsConfig.GetCertificate = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {

		// Load from memory
		// TODO unload untouched certificates every x minutes
		if myCert, ok := certMap[clientHello.ServerName]; ok {
			myCert.touchedAt = time.Now()
			return myCert.cert, nil
		}

		privkeyPath := filepath.Join(*certsPath, clientHello.ServerName, "privkey.pem")
		certPath := filepath.Join(*certsPath, clientHello.ServerName, "fullchain.pem")

		loadCert := func() *tls.Certificate {
			// TODO handle race condition (ask Matt)
			// the transaction is idempotent, however, so it shouldn't matter
			if _, err := os.Stat(privkeyPath); err == nil {
				fmt.Printf("Loading Certificates %s/%s/{privkey.pem,fullchain.pem}\n\n", *certsPath, clientHello.ServerName)
				cert, err := tls.LoadX509KeyPair(certPath, privkeyPath)
				if nil != err {
					return &cert
				}
				return nil
			}

			return nil
		}

		if cert := loadCert(); nil != cert {
			certMap[clientHello.ServerName] = myCert{
				cert:      cert,
				touchedAt: time.Now(),
			}
			return cert, nil
		}

		// TODO try to get cert via letsencrypt python client
		// TODO check for a hosting directory before attempting this
		/*
			cmd := exec.Command(
				"./venv/bin/letsencrypt",
				"--text",
				"--agree-eula",
				"--email", "coolaj86@gmail.com",
				"--authenticator", "standalone",
				"--domains", "www.example.com", "example.com",
				"--dvsni-port", "65443",
				"auth",
			)
			err := cmd.Run()
			if nil != err {
				if cert := loadCert(); nil != cert {
					return cert, nil
				}
			}
		*/

		fmt.Fprintf(os.Stderr, "Failed to load certificates for %q.\n", clientHello.ServerName)
		fmt.Fprintf(os.Stderr, "\tTried %s/{privkey.pem,fullchain.pem}\n", filepath.Join(*certsPath, clientHello.ServerName))
		//fmt.Fprintf(os.Stderr, "\tand letsencrypt api\n")
		fmt.Fprintf(os.Stderr, "\n")
		// TODO how to prevent attack and still enable retry?
		// perhaps check DNS and hosting directory, wait 5 minutes?
		certMap[clientHello.ServerName] = myCert{
			cert:      &defaultCert,
			touchedAt: time.Now(),
		}
		return &defaultCert, nil
	}
	tlsListener := tls.NewListener(conn, tlsConfig)

	server := &http.Server{
		Addr:    addr,
		Handler: &myHandler{},
	}
	fmt.Printf("Listening on https://%s:%d\n\n", host, *port)
	server.Serve(tlsListener)
}
<commit_msg>correct comment about arg<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

func usage() {
	fmt.Fprintf(os.Stderr, "\nusage: go run serve.go [optional flags]\n")
	flag.PrintDefaults()
	fmt.Println()

	os.Exit(2)
}

type myHandler struct {
	certMap map[string]tls.Certificate
}

type myCert struct {
	cert      *tls.Certificate
	touchedAt time.Time
}

func (m *myHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Print debug info
	fmt.Println(r.Host)
	fmt.Println(r.Method)
	fmt.Println(r.RequestURI)
	fmt.Println(r.URL) // has many keys, such as Query
	for k, v := range r.Header {
		fmt.Println(k, v)
	}

	fmt.Println(r.Body)
	fmt.Println()
	fmt.Println()

	// End the request
	// TODO serve from hosting directory
	fmt.Fprintf(w, "Hi there, %s %q? Wow!\n\nWith Love,\n\t%s", r.Method, r.URL.Path[1:], r.Host)
}

func main() {
	flag.Usage = usage

	port := flag.Uint("port", 443, "https port")
	certsPath := flag.String("letsencrypt-path", "/etc/letsencrypt/live", "path at which an 'xyz.example.com' containing 'fullchain.pem' and 'privkey.pem' can be found")
	defaultHost := flag.String("default-hostname", "localhost.daplie.com", "the default folder to find certificates to use when no matches are found")

	flag.Parse()

	host := strings.ToLower(*defaultHost)
	// See https://groups.google.com/a/letsencrypt.org/forum/#!topic/ca-dev/l1Dd6jzWeu8
	/*
		if strings.HasPrefix("www.", host) {
			fmt.Println("TODO: 'www.' prefixed certs should be obtained for every 'example.com' domain.")
		}
		host = strings.TrimPrefix("www.", host)
	*/

	fmt.Printf("Loading Certificates %s/%s/{privkey.pem,fullchain.pem}\n", *certsPath, *defaultHost)
	privkeyPath := filepath.Join(*certsPath, *defaultHost, "privkey.pem")
	certPath := filepath.Join(*certsPath, *defaultHost, "fullchain.pem")
	defaultCert, err := tls.LoadX509KeyPair(certPath, privkeyPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't load default certificates: %s\n", err)
		os.Exit(1)
	}

	addr := ":" + strconv.Itoa(int(*port))

	conn, err := net.Listen("tcp", addr)
	if nil != err {
		fmt.Fprintf(os.Stderr, "Couldn't bind to TCP socket %q: %s\n", addr, err)
		os.Exit(1)
	}

	certMap := make(map[string]myCert)
	tlsConfig := new(tls.Config)
	tlsConfig.Certificates = []tls.Certificate{defaultCert}
	tlsConfig.GetCertificate = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {

		// Load from memory
		// TODO unload untouched certificates every x minutes
		if myCert, ok := certMap[clientHello.ServerName]; ok {
			myCert.touchedAt = time.Now()
			return myCert.cert, nil
		}

		privkeyPath := filepath.Join(*certsPath, clientHello.ServerName, "privkey.pem")
		certPath := filepath.Join(*certsPath, clientHello.ServerName, "fullchain.pem")

		loadCert := func() *tls.Certificate {
			// TODO handle race condition (ask Matt)
			// the transaction is idempotent, however, so it shouldn't matter
			if _, err := os.Stat(privkeyPath); err == nil {
				fmt.Printf("Loading Certificates %s/%s/{privkey.pem,fullchain.pem}\n\n", *certsPath, clientHello.ServerName)
				cert, err := tls.LoadX509KeyPair(certPath, privkeyPath)
				if nil != err {
					return &cert
				}
				return nil
			}

			return nil
		}

		if cert := loadCert(); nil != cert {
			certMap[clientHello.ServerName] = myCert{
				cert:      cert,
				touchedAt: time.Now(),
			}
			return cert, nil
		}

		// TODO try to get cert via letsencrypt python client
		// TODO check for a hosting directory before attempting this
		/*
			cmd := exec.Command(
				"./venv/bin/letsencrypt",
				"--text",
				"--agree-eula",
				"--email", "coolaj86@gmail.com",
				"--authenticator", "standalone",
				"--domains", "www.example.com",
				"--domains", "example.com",
				"--dvsni-port", "65443",
				"auth",
			)
			err := cmd.Run()
			if nil != err {
				if cert := loadCert(); nil != cert {
					return cert, nil
				}
			}
		*/

		fmt.Fprintf(os.Stderr, "Failed to load certificates for %q.\n", clientHello.ServerName)
		fmt.Fprintf(os.Stderr, "\tTried %s/{privkey.pem,fullchain.pem}\n", filepath.Join(*certsPath, clientHello.ServerName))
		//fmt.Fprintf(os.Stderr, "\tand letsencrypt api\n")
		fmt.Fprintf(os.Stderr, "\n")
		// TODO how to prevent attack and still enable retry?
		// perhaps check DNS and hosting directory, wait 5 minutes?
		certMap[clientHello.ServerName] = myCert{
			cert:      &defaultCert,
			touchedAt: time.Now(),
		}
		return &defaultCert, nil
	}
	tlsListener := tls.NewListener(conn, tlsConfig)

	server := &http.Server{
		Addr:    addr,
		Handler: &myHandler{},
	}
	fmt.Printf("Listening on https://%s:%d\n\n", host, *port)
	server.Serve(tlsListener)
}
<|endoftext|>


<commit_before>package arn

import "github.com/aerogo/nano"

// Notification represents a user-associated notification.
type Notification struct {
	ID      string `json:"id"`
	UserID  string `json:"userId"`
	Created string `json:"created"`
	Seen    string `json:"seen"`
	PushNotification
}

// User retrieves the user the notification was sent to.
func (notification *Notification) User() *User {
	user, _ := GetUser(notification.UserID)
	return user
}

// NewNotification creates a new notification.
func NewNotification(userID string, pushNotification *PushNotification) *Notification {
	return &Notification{
		ID:               GenerateID("Notification"),
		UserID:           userID,
		Created:          DateTimeUTC(),
		Seen:             "",
		PushNotification: *pushNotification,
	}
}

// GetNotification ...
func GetNotification(id string) (*Notification, error) {
	obj, err := DB.Get("Notification", id)

	if err != nil {
		return nil, err
	}

	return obj.(*Notification), nil
}

// StreamNotifications returns a stream of all notifications.
func StreamNotifications() chan *Notification {
	channel := make(chan *Notification, nano.ChannelBufferSize)

	go func() {
		for obj := range DB.All("Notification") {
			channel <- obj.(*Notification)
		}

		close(channel)
	}()

	return channel
}

// AllNotifications returns a slice of all notifications.
func AllNotifications() ([]*Notification, error) {
	var all []*Notification

	for obj := range StreamNotifications() {
		all = append(all, obj)
	}

	return all, nil
}
<commit_msg>Added notification.CreatedTime<commit_after>package arn

import (
	"time"

	"github.com/aerogo/nano"
)

// Notification represents a user-associated notification.
type Notification struct {
	ID      string `json:"id"`
	UserID  string `json:"userId"`
	Created string `json:"created"`
	Seen    string `json:"seen"`
	PushNotification
}

// User retrieves the user the notification was sent to.
func (notification *Notification) User() *User {
	user, _ := GetUser(notification.UserID)
	return user
}

// CreatedTime returns the created date as a time object.
func (notification *Notification) CreatedTime() time.Time {
	t, _ := time.Parse(time.RFC3339, notification.Created)
	return t
}

// NewNotification creates a new notification.
func NewNotification(userID string, pushNotification *PushNotification) *Notification {
	return &Notification{
		ID:               GenerateID("Notification"),
		UserID:           userID,
		Created:          DateTimeUTC(),
		Seen:             "",
		PushNotification: *pushNotification,
	}
}

// GetNotification ...
func GetNotification(id string) (*Notification, error) {
	obj, err := DB.Get("Notification", id)

	if err != nil {
		return nil, err
	}

	return obj.(*Notification), nil
}

// StreamNotifications returns a stream of all notifications.
func StreamNotifications() chan *Notification {
	channel := make(chan *Notification, nano.ChannelBufferSize)

	go func() {
		for obj := range DB.All("Notification") {
			channel <- obj.(*Notification)
		}

		close(channel)
	}()

	return channel
}

// AllNotifications returns a slice of all notifications.
func AllNotifications() ([]*Notification, error) {
	var all []*Notification

	for obj := range StreamNotifications() {
		all = append(all, obj)
	}

	return all, nil
}
<|endoftext|>


<commit_before>package eventsocket

import (
	"time"

	log "github.com/dronemill/eventsocket/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/dronemill/eventsocket/Godeps/_workspace/src/github.com/gorilla/websocket"
)

const (
	// time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// time allowed to read the next pong message from the peer.
	pongWait = 15 * time.Second

	// send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
)

// maximum message size allowed from peer
var defaultMaxMessageSize = int64(512)
var defaultReadDeadline = 30 * time.Second

type wsConnection struct {
	// The websocket connection.
	ws *websocket.Conn

	// Buffered channel of outbound messages.
	send chan Message

	// Receive messages from this client
	recv chan Message

	// has this wsConnection been closed
	closed bool

	// the maximum message size allowed
	maxMessageSize int64

	// the socket Read Deadline
	readDeadline time.Duration
}

func newWsConnection(ws *websocket.Conn) (*wsConnection, error) {
	log.Info("Creating new wsConn")

	wsc := &wsConnection{
		send:           make(chan Message, 256),
		recv:           make(chan Message, 256),
		ws:             ws,
		closed:         false,
		maxMessageSize: defaultMaxMessageSize,
		readDeadline:   defaultReadDeadline,
	}

	return wsc, nil
}

// override the default readlimit of 512
func (wsc *wsConnection) SetMaxMessageSize(limit int64) {
	wsc.maxMessageSize = limit
}

// override the default readDeadline
func (wsc *wsConnection) SetReadDeadline(t time.Duration) {
	wsc.readDeadline = t
}

func (wsc *wsConnection) pump() {
	go wsc.writePump()
	wsc.readPump()
}

// readPump pumps messages from the websocket connection to the hub.
func (wsc *wsConnection) readPump() {
	defer func() {
		h.unregister <- wsc
	}()

	wsc.ws.SetReadLimit(wsc.maxMessageSize)
	wsc.ws.SetReadDeadline(time.Now().Add(wsc.readDeadline))
	wsc.ws.SetPongHandler(func(string) error { wsc.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		m := &Message{}
		if wsc.ws.ReadJSON(m) != nil {
			break
		}

		wsc.recv <- *m
	}
}

// write writes a message with the given message type and payload.
func (wsc *wsConnection) write(mt int, payload []byte) error {
	wsc.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return wsc.ws.WriteMessage(mt, payload)
}

// writeJSON writes a json message
func (wsc *wsConnection) writeJSON(message interface{}) error {
	wsc.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return wsc.ws.WriteJSON(message)
}

// writePump pumps messages from the hub to the websocket connection.
func (wsc *wsConnection) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		h.unregister <- wsc
	}()
	for {
		select {
		case message, ok := <-wsc.send:
			if !ok {
				wsc.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := wsc.writeJSON(message); err != nil {
				return
			}
		case <-ticker.C:
			if err := wsc.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (wsc *wsConnection) Close() {
	// if we have already closed the channels, then get out
	if wsc.closed {
		return
	}
	wsc.closed = true

	// blose stuff
	close(wsc.send)
	close(wsc.recv)
	wsc.ws.Close()
}
<commit_msg>Added ping logging<commit_after>package eventsocket

import (
	"time"

	log "github.com/dronemill/eventsocket/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/dronemill/eventsocket/Godeps/_workspace/src/github.com/gorilla/websocket"
)

const (
	// time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// time allowed to read the next pong message from the peer.
	pongWait = 15 * time.Second

	// send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
)

// maximum message size allowed from peer
var defaultMaxMessageSize = int64(512)
var defaultReadDeadline = 30 * time.Second

type wsConnection struct {
	// The ID of the websocket
	id string

	// The websocket connection.
	ws *websocket.Conn

	// Buffered channel of outbound messages.
	send chan Message

	// Receive messages from this client
	recv chan Message

	// has this wsConnection been closed
	closed bool

	// the maximum message size allowed
	maxMessageSize int64

	// the socket Read Deadline
	readDeadline time.Duration
}

func newWsConnection(ws *websocket.Conn) (*wsConnection, error) {
	id := <-uuidBuilder
	log.WithField("websocketID", id.String()).Info("Creating new wsConn")

	wsc := &wsConnection{
		id:             id.String(),
		send:           make(chan Message, 256),
		recv:           make(chan Message, 256),
		ws:             ws,
		closed:         false,
		maxMessageSize: defaultMaxMessageSize,
		readDeadline:   defaultReadDeadline,
	}

	return wsc, nil
}

// override the default readlimit of 512
func (wsc *wsConnection) SetMaxMessageSize(limit int64) {
	wsc.maxMessageSize = limit
}

// override the default readDeadline
func (wsc *wsConnection) SetReadDeadline(t time.Duration) {
	wsc.readDeadline = t
}

func (wsc *wsConnection) pump() {
	go wsc.writePump()
	wsc.readPump()
}

// readPump pumps messages from the websocket connection to the hub.
func (wsc *wsConnection) readPump() {
	defer func() {
		h.unregister <- wsc
	}()

	wsc.ws.SetReadLimit(wsc.maxMessageSize)
	wsc.ws.SetReadDeadline(time.Now().Add(pongWait))
	wsc.ws.SetPongHandler(func(string) error { wsc.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		m := &Message{}
		if wsc.ws.ReadJSON(m) != nil {
			break
		}

		wsc.recv <- *m
	}
}

// write writes a message with the given message type and payload.
func (wsc *wsConnection) write(mt int, payload []byte) error {
	wsc.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return wsc.ws.WriteMessage(mt, payload)
}

// writeJSON writes a json message
func (wsc *wsConnection) writeJSON(message interface{}) error {
	wsc.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return wsc.ws.WriteJSON(message)
}

// writePump pumps messages from the hub to the websocket connection.
func (wsc *wsConnection) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		h.unregister <- wsc
	}()
	for {
		select {
		case message, ok := <-wsc.send:
			if !ok {
				wsc.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := wsc.writeJSON(message); err != nil {
				return
			}
		case <-ticker.C:
			log.WithField("websocketID", wsc.id).Info("Pinging wsConn")
			if err := wsc.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (wsc *wsConnection) Close() {
	// if we have already closed the channels, then get out
	if wsc.closed {
		return
	}
	wsc.closed = true

	// blose stuff
	close(wsc.send)
	close(wsc.recv)
	wsc.ws.Close()
}
<|endoftext|>


<commit_before>// Package trie implements a general purpose Path-Trie.
package trie

import (
	"os"
	"strings"
)

// Node represents a single node in a Trie, but it can be used as a whole
// (sub-)Trie through the Trie interface. A node value of `nil` is a perfectly
// valid trie. Node is suitable for embedding it into other structs.
type Node struct {
	// Pointer to parent node or nil
	Parent *Node

	// Basename to child-nodes
	Children map[string]*Node

	// Basename of the node's Path
	Name string

	// Number of explicitly added children of this node.
	// (1 for leaf nodes)
	Length int64

	// Depth of the node. The root is at depth 0.
	Depth uint16
}

// Trie represents the required methods for accessing a directory structure.
type Trie interface {
	// Root returns the uppermost node of the trie.
	Root() *Node

	// Insert adds a new node in the trie at string. If the node already exists,
	// nothing changes. This operation costs O(log(n)). The newly created or
	// existant node is returned.
	Insert(path string) *Node

	// Lookup searches for a node references by a path.
	Lookup(path string) *Node

	// Remove removes the node at path and all of it's children.
	// The parent of the removed node is returned, which might be nil.
	Remove() *Node

	// Len returns the current number of elements in the trie.
	// This counts only explicitly inserted Nodes.
	Len() int64
}

// SplitPath splits the path according to os.PathSeparator,
// but omits a leading empty name on /unix/paths
func SplitPath(path string) []string {
	names := strings.Split(path, string(os.PathSeparator))
	if len(names) > 0 && names[0] == "" {
		return names[1:]
	}

	return names
}

// NewTrie returns a trie with the root element pre-inserted.
// Note that `nil` is a perfectly valid, but empty trie.
func NewTrie() *Node {
	return &Node{}
}

// Root returns the root node of the trie.
func (n *Node) Root() *Node {
	if n != nil && n.Parent != nil {
		return n.Parent.Root()
	}
	return n
}

// Insert adds a node into the trie at `path`
func (n *Node) Insert(path string) *Node {
	curr := n
	if curr == nil {
		curr = NewTrie()
	}

	wasAdded := false

	for depth, name := range SplitPath(path) {
		if curr.Children == nil {
			curr.Children = make(map[string]*Node)
		}
		child, ok := curr.Children[name]
		if !ok {
			child = &Node{
				Parent: curr,
				Name:   name,
				Depth:  uint16(depth + 1),
			}
			curr.Children[name] = child
			wasAdded = true
		}
		curr = child
	}

	if wasAdded && curr != nil {
		curr.Up(func(parent *Node) {
			parent.Length++
		})
	}
	return curr
}

// Lookup searches a Node by it's absolute path.
func (n *Node) Lookup(path string) *Node {
	curr := n
	if n == nil {
		return nil
	}

	if path == "/" {
		return n.Root()
	}

	for _, name := range SplitPath(path) {
		child, ok := curr.Children[name]
		if !ok {
			return nil
		}
		curr = child
	}
	return curr
}

// Remove removes the receiver and all of it's children.
// The removed node's parent is returned.
func (n *Node) Remove() *Node {
	if n == nil {
		return nil
	}

	// Adjusts the parent's length:
	length := n.Length
	n.Up(func(parent *Node) {
		parent.Length -= length
	})

	// Removes link to self:
	if n.Parent != nil {
		delete(n.Parent.Children, n.Name)
	}

	// Make children garbage collectable:
	parent := n.Parent
	n.Walk(true, func(child *Node) {
		child.Children = nil
		child.Parent = nil
	})
	return parent
}

// Walk iterates over all (including intermediate )nodes in the trie.
// Depending on dfs the nodes are visited in depth-first or breadth-first.
// The supplied callback is called for each visited node.
func (n *Node) Walk(dfs bool, visit func(*Node)) {
	if n == nil {
		return
	}

	if !dfs {
		visit(n)
	}

	if n.Children != nil {
		for _, child := range n.Children {
			child.Walk(dfs, visit)
		}
	}

	if dfs {
		visit(n)
	}
}

// Up walks from the receiving node to the root node,
// calling `visit` on each node on it's way.
func (n *Node) Up(visit func(*Node)) {
	if n != nil {
		visit(n)
		n.Parent.Up(visit)
	}
}

// Len returns the number of explicitly inserted elements in the trie.
func (n *Node) Len() int64 {
	if n == nil {
		return 0
	}
	return n.Length
}

// Path returns a full absolute path from the receiver
// to the root node of the trie.
func (n *Node) Path() string {
	if n == nil {
		return ""
	}

	// Determinate the path length.
	pathLen := 0
	n.Up(func(parent *Node) {
		pathLen++
	})

	pathChunks := make([]string, pathLen+1)

	i := pathLen
	n.Up(func(parent *Node) {
		if i >= 0 {
			pathChunks[i] = parent.Name
			i--
		}
	})

	return buildPath(pathChunks)
}

// String returns the absolute path of the node.
func (n *Node) String() string {
	if n == nil {
		return "<nil>"
	}
	return n.Path()
}
<commit_msg>pathricia.go: Path determine length corrected by one.<commit_after>// Package trie implements a general purpose Path-Trie.
package trie

import (
	"os"
	"strings"
)

// Node represents a single node in a Trie, but it can be used as a whole
// (sub-)Trie through the Trie interface. A node value of `nil` is a perfectly
// valid trie. Node is suitable for embedding it into other structs.
type Node struct {
	// Pointer to parent node or nil
	Parent *Node

	// Basename to child-nodes
	Children map[string]*Node

	// Basename of the node's Path
	Name string

	// Number of explicitly added children of this node.
	// (1 for leaf nodes)
	Length int64

	// Depth of the node. The root is at depth 0.
	Depth uint16
}

// Trie represents the required methods for accessing a directory structure.
type Trie interface {
	// Root returns the uppermost node of the trie.
	Root() *Node

	// Insert adds a new node in the trie at string. If the node already exists,
	// nothing changes. This operation costs O(log(n)). The newly created or
	// existant node is returned.
	Insert(path string) *Node

	// Lookup searches for a node references by a path.
	Lookup(path string) *Node

	// Remove removes the node at path and all of it's children.
	// The parent of the removed node is returned, which might be nil.
	Remove() *Node

	// Len returns the current number of elements in the trie.
	// This counts only explicitly inserted Nodes.
	Len() int64
}

// SplitPath splits the path according to os.PathSeparator,
// but omits a leading empty name on /unix/paths
func SplitPath(path string) []string {
	names := strings.Split(path, string(os.PathSeparator))
	if len(names) > 0 && names[0] == "" {
		return names[1:]
	}

	return names
}

// NewTrie returns a trie with the root element pre-inserted.
// Note that `nil` is a perfectly valid, but empty trie.
func NewTrie() *Node {
	return &Node{}
}

// Root returns the root node of the trie.
func (n *Node) Root() *Node {
	if n != nil && n.Parent != nil {
		return n.Parent.Root()
	}
	return n
}

// Insert adds a node into the trie at `path`
func (n *Node) Insert(path string) *Node {
	curr := n
	if curr == nil {
		curr = NewTrie()
	}

	wasAdded := false

	for depth, name := range SplitPath(path) {
		if curr.Children == nil {
			curr.Children = make(map[string]*Node)
		}
		child, ok := curr.Children[name]
		if !ok {
			child = &Node{
				Parent: curr,
				Name:   name,
				Depth:  uint16(depth + 1),
			}
			curr.Children[name] = child
			wasAdded = true
		}
		curr = child
	}

	if wasAdded && curr != nil {
		curr.Up(func(parent *Node) {
			parent.Length++
		})
	}
	return curr
}

// Lookup searches a Node by it's absolute path.
func (n *Node) Lookup(path string) *Node {
	curr := n
	if n == nil {
		return nil
	}

	if path == "/" {
		return n.Root()
	}

	for _, name := range SplitPath(path) {
		child, ok := curr.Children[name]
		if !ok {
			return nil
		}
		curr = child
	}
	return curr
}

// Remove removes the receiver and all of it's children.
// The removed node's parent is returned.
func (n *Node) Remove() *Node {
	if n == nil {
		return nil
	}

	// Adjusts the parent's length:
	length := n.Length
	n.Up(func(parent *Node) {
		parent.Length -= length
	})

	// Removes link to self:
	if n.Parent != nil {
		delete(n.Parent.Children, n.Name)
	}

	// Make children garbage collectable:
	parent := n.Parent
	n.Walk(true, func(child *Node) {
		child.Children = nil
		child.Parent = nil
	})
	return parent
}

// Walk iterates over all (including intermediate )nodes in the trie.
// Depending on dfs the nodes are visited in depth-first or breadth-first.
// The supplied callback is called for each visited node.
func (n *Node) Walk(dfs bool, visit func(*Node)) {
	if n == nil {
		return
	}

	if !dfs {
		visit(n)
	}

	if n.Children != nil {
		for _, child := range n.Children {
			child.Walk(dfs, visit)
		}
	}

	if dfs {
		visit(n)
	}
}

// Up walks from the receiving node to the root node,
// calling `visit` on each node on it's way.
func (n *Node) Up(visit func(*Node)) {
	if n != nil {
		visit(n)
		n.Parent.Up(visit)
	}
}

// Len returns the number of explicitly inserted elements in the trie.
func (n *Node) Len() int64 {
	if n == nil {
		return 0
	}
	return n.Length
}

// Path returns a full absolute path from the receiver
// to the root node of the trie.
func (n *Node) Path() string {
	if n == nil {
		return ""
	}

	// Determine the path length,
	// the path e.g. /home/a/b/c has a length of four.
	pathLen := 0
	n.Up(func(parent *Node) {
		pathLen++
	})

	pathChunks := make([]string, pathLen)

	i := pathLen
	n.Up(func(parent *Node) {
		if i > 0 {
			pathChunks[i-1] = parent.Name
			i--
		}
	})

	return buildPath(pathChunks)
}

// String returns the absolute path of the node.
func (n *Node) String() string {
	if n == nil {
		return "<nil>"
	}
	return n.Path()
}
<|endoftext|>


<commit_before>/*
 *   Copyright 2016 Rémy MATHIEU
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

package content

import (
	"strings"
	"github.com/PuerkitoBio/goquery"
	"log"
)

type RedditContent struct {
	Url string
}

func (reddit RedditContent) callAPI() ([]Content, error) {
	resp, err := getWebserviceResponse(reddit.Url)
	if( err != nil ) {
		log.Println("Error while calling url: "+reddit.Url)
		return nil, err
	}

	doc, err := goquery.NewDocumentFromResponse(resp)
	if err != nil {
		log.Println("Error while calling API")
		return nil, err
	}

	rv := make([]Content, 0)

	doc.Find("div.search-result,div.entry").Each(func(i int, selec *goquery.Selection) {
		// ignore sticky posts
		if selec.HasClass("stickied") {
			return
		}

		if len(rv) > 20 {
			return
		}

		t := selec.Find("a.search-title,a.title")
		title := t.First().Text()

		// Limit size of content
		if( len(title) + urlLength > 140 ) {
			title = title[0:139-urlLength] + "…"
		}

		l := selec.Find("a.search-link,a.link")
		externalLink, _ := l.First().Attr("href")

		// self posts
		if strings.HasPrefix(externalLink, "/r/") {
 			externalLink = "https://reddit.com" + externalLink
 		}

		rv = append(rv, Content{
			Text: title,
			Url:  externalLink,
		})
	})

	return rv, nil
}
<commit_msg>all: gofmt<commit_after>/*
 *   Copyright 2016 Rémy MATHIEU
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

package content

import (
	"github.com/PuerkitoBio/goquery"
	"log"
	"strings"
)

type RedditContent struct {
	Url string
}

func (reddit RedditContent) callAPI() ([]Content, error) {
	resp, err := getWebserviceResponse(reddit.Url)
	if err != nil {
		log.Println("Error while calling url: " + reddit.Url)
		return nil, err
	}

	doc, err := goquery.NewDocumentFromResponse(resp)
	if err != nil {
		log.Println("Error while calling API")
		return nil, err
	}

	rv := make([]Content, 0)

	doc.Find("div.search-result,div.entry").Each(func(i int, selec *goquery.Selection) {
		// ignore sticky posts
		if selec.HasClass("stickied") {
			return
		}

		if len(rv) > 20 {
			return
		}

		t := selec.Find("a.search-title,a.title")
		title := t.First().Text()

		// Limit size of content
		if len(title)+urlLength > 140 {
			title = title[0:139-urlLength] + "…"
		}

		l := selec.Find("a.search-link,a.link")
		externalLink, _ := l.First().Attr("href")

		// self posts
		if strings.HasPrefix(externalLink, "/r/") {
			externalLink = "https://reddit.com" + externalLink
		}

		rv = append(rv, Content{
			Text: title,
			Url:  externalLink,
		})
	})

	return rv, nil
}
<|endoftext|>


<commit_before>
package main
import (
	"fmt"
	"os"
	"encoding/json"
	"io/ioutil"
	"log"
)

//UserInfo stores information of user
type UserInfo struct {
	First string
	Last  string
}

//Check checks for errors
func check(e error) {
	if e != nil {
		panic(e)
	}
}

//printHelp opens project page (http://hasit.github.io/licensethis/) in the default browser.
func printHelp() {
	helptext:=`
# licensethis
Add an OSS license to your project.

### SYNOPSIS

licensethis <option> <license-name>
licensethis help
licensethis config
licensethis list
licensethis info license-name
licensethis generate license-name

DESCRIPTION

licensethis lets you add an OSS license to your project with ease.

OPTIONS

These options let you do a variety of things.

  * help
	Open project page in the default browser. This provides more information on usage of licensethis with detailed information about each option.
  * config
	Configure user's full name for generating license files. User will be prompted to enter their full name for the first time. This information is stored in a JSON file at *~/.config/licensethis/user.json*.
  * list
	List all OSS licenses available on *http://choosealicense.com/*. This only prints a list of licenses with no more information. This could serve as a quick preview of all available licenses to choose from.
  * info license-name
	Get more information for <license-name>. This prints a synopsis of the license, tags (required, permitted, forbidden) associated to the license and a link to *http://choosealicense.com/licenses/<license-name>*.
  * generate license-name
	Generate <license-name> license for your project. This generates a LICENSE.txt file in your current directory with YEAR and FULLNAME already filled in. It picks up the name from config file automatically.

## CONTRIBUTORS
* Hasit Mistry
	Personal: [hasit.github.io](http://hasit.github.io/)
	Github: [hasit](https://github.com/hasit)
* Anuj Deshpande
	Personal: [anujdeshpande.com](http://www.anujdeshpande.com/)
	Github: [anujdeshpande](https://github.com/anujdeshpande/)\n`
	fmt.Printf("%v",helptext)
}

func userConfig() {
	first := os.Getenv("FIRSTNAME")
	last := os.Getenv("LASTNAME")

	//if first != nil; last != nil {
	fmt.Printf("First: %v\n", first)
	fmt.Printf("Last: %v\n", last)
	//}
}

func info(licensename []string) {
	fmt.Printf("%v\n", licensename)

	file, e := ioutil.ReadFile("licenses.json")
	if e != nil {
		log.Fatal(e)
	}
	var jsontype map[string][]map[string]interface{}
	err := json.Unmarshal (file,&jsontype)
	if err!=nil {
		log.Fatal(err)
	}
	//fmt.Printf("Results: %v \n", jsontype["licenses"])
	for i := range jsontype["licenses"] {
		item := jsontype["licenses"][i]
		fmt.Printf("%v\n",item["name"])
	}
}

//parseArgs parses command line arguments and calls appropriate functions.
func parseArgs(args []string) {
	if len(args) != 0 {
		switch args[0] {
		case "help":
			printHelp()
		case "config":
			userConfig()
		case "info":
			info(args[1:])
		case "list":
			fmt.Println("list")
		case "generate":
			fmt.Println("generate")
		default:
			fmt.Println("Incorrect usage!")
			fmt.Println("Type `licensethis help` for help on proper usage.")
		}
	} else {
		fmt.Println("Incorrect usage!")
		fmt.Println("Type `licensethis help` for help on proper usage.")
	}
}

func main() {
	args := os.Args[1:]

	parseArgs(args)
}
<commit_msg>Read license text from a json : Part III<commit_after>
package main
import (
	"fmt"
	"os"
	"encoding/json"
	"io/ioutil"
	"log"
)

//UserInfo stores information of user
type UserInfo struct {
	First string
	Last  string
}

//Check checks for errors
func check(e error) {
	if e != nil {
		panic(e)
	}
}

//printHelp opens project page (http://hasit.github.io/licensethis/) in the default browser.
func printHelp() {
	helptext:=`
# licensethis
Add an OSS license to your project.

### SYNOPSIS

licensethis <option> <license-name>
licensethis help
licensethis config
licensethis list
licensethis info license-name
licensethis generate license-name

DESCRIPTION

licensethis lets you add an OSS license to your project with ease.

OPTIONS

These options let you do a variety of things.

  * help
	Open project page in the default browser. This provides more information on usage of licensethis with detailed information about each option.
  * config
	Configure user's full name for generating license files. User will be prompted to enter their full name for the first time. This information is stored in a JSON file at *~/.config/licensethis/user.json*.
  * list
	List all OSS licenses available on *http://choosealicense.com/*. This only prints a list of licenses with no more information. This could serve as a quick preview of all available licenses to choose from.
  * info license-name
	Get more information for <license-name>. This prints a synopsis of the license, tags (required, permitted, forbidden) associated to the license and a link to *http://choosealicense.com/licenses/<license-name>*.
  * generate license-name
	Generate <license-name> license for your project. This generates a LICENSE.txt file in your current directory with YEAR and FULLNAME already filled in. It picks up the name from config file automatically.

## CONTRIBUTORS
* Hasit Mistry
	Personal: [hasit.github.io](http://hasit.github.io/)
	Github: [hasit](https://github.com/hasit)
* Anuj Deshpande
	Personal: [anujdeshpande.com](http://www.anujdeshpande.com/)
	Github: [anujdeshpande](https://github.com/anujdeshpande/)\n`
	fmt.Printf("%v",helptext)
}

func userConfig() {
	first := os.Getenv("FIRSTNAME")
	last := os.Getenv("LASTNAME")

	//if first != nil; last != nil {
	fmt.Printf("First: %v\n", first)
	fmt.Printf("Last: %v\n", last)
	//}
}

func info(licensename []string) {
	//fmt.Printf("%v\n", licensename[0])

	file, e := ioutil.ReadFile("licenses.json")
	if e != nil {
		log.Fatal(e)
	}
	var jsontype map[string][]map[string]interface{}
	err := json.Unmarshal (file,&jsontype)
	if err!=nil {
		log.Fatal(err)
	}
	//fmt.Printf("Results: %v \n", jsontype["licenses"])
	for i := range jsontype["licenses"] {
		item := jsontype["licenses"][i]
		//fmt.Printf("%v\n",item["name"])
		if item["name"]==licensename[0] {
			fmt.Printf("%v\n%v\n%v", item["name"], item["version"],item["text"])
			}
	}

}

//parseArgs parses command line arguments and calls appropriate functions.
func parseArgs(args []string) {
	if len(args) != 0 {
		switch args[0] {
		case "help":
			printHelp()
		case "config":
			userConfig()
		case "info":
			info(args[1:])
		case "list":
			fmt.Println("list")
		case "generate":
			fmt.Println("generate")
		default:
			fmt.Println("Incorrect usage!")
			fmt.Println("Type `licensethis help` for help on proper usage.")
		}
	} else {
		fmt.Println("Incorrect usage!")
		fmt.Println("Type `licensethis help` for help on proper usage.")
	}
}

func main() {
	args := os.Args[1:]

	parseArgs(args)
}
<|endoftext|>


<commit_before>package lifted

import (
	"fmt"
	"bytes"
)

func (d *Domain) String() string {
	buf := bytes.NewBuffer(make([]byte, 0, 100))

	buf.WriteString("&Domain{")
	fmt.Fprintf(buf, "Name:%s,\n", d.Name)
	fmt.Fprintf(buf, "Requirements:%+v\n", d.Requirements)
	fmt.Fprintf(buf, "Types:%v\n", d.Types)
	fmt.Fprintf(buf, "Constants:%v\n", d.Constants)
	fmt.Fprintf(buf, "Predicates:%+v\n", d.Predicates)

	buf.WriteString("Actions:[\n")
	for i, a := range d.Actions {
		buf.WriteString(a.String())
		if i < len(d.Actions)-1 {
			buf.WriteString("\n\n")
		}
	}
	buf.WriteString("],\n}")

	return buf.String()
}

func (p Predicate) String() string {
	return fmt.Sprintf("Predicate{Name:%v, Parameters:%v}",
		p.Name, p.Parameters)
}

func (a Action) String() string {
	buf := bytes.NewBuffer(make([]byte, 0, 100))

	buf.WriteString("Action{")
	fmt.Fprintf(buf, "Name:%s", a.Name)
	fmt.Fprintf(buf, "\nParameters:%+v", a.Parameters)
	fmt.Fprintf(buf, "\nPrecondition:%+v", a.Precondition)
	fmt.Fprintf(buf, "\nEffect:%+v", a.Effect)
	buf.WriteString("}")

	return buf.String()
}

func (p *Problem) String() string {
	buf := bytes.NewBuffer(make([]byte, 0, 100))

	buf.WriteString("Problem{")
	fmt.Fprintf(buf, "Name:%s\n", p.Name)
	fmt.Fprintf(buf, "Domain:%s\n", p.Domain)
	fmt.Fprintf(buf, "Requirements:%v\n", p.Requirements)
	fmt.Fprintf(buf, "Objects:%v\n", p.Objects)
	fmt.Fprintf(buf, "Init:%+v\n", p.Init)
	fmt.Fprintf(buf, "Goal:%+v\n", p.Goal)
	fmt.Fprintf(buf, "Metric:%+v\n", p.Metric)
	buf.WriteByte('}')

	return buf.String()
}

func (k TermKind) String() string {
	switch k {
	case TermVariable:
		return "TermVariable"
	case TermConstant:
		return "TermConstant"
	}

	return fmt.Sprintf("%d", int(k))
}

func (t Term) String() string {
	return fmt.Sprintf("Term{Kind:%v, Name:%v}",
		t.Kind, t.Name)
}

func (n Name) String() string {
	if n.Num < 0 {
		return fmt.Sprintf("{%s}", n.Str)
	}
	return fmt.Sprintf("{%s, %d}", n.Str, n.Num)
}

func (lit *LiteralNode) String() string {
	return fmt.Sprintf("Literal{Positive:%t, Name:%v, Parameters:%v}",
		lit.Positive, lit.Name, lit.Parameters)
}

func (e *BinaryNode) String() string {
	return fmt.Sprintf("Left:%v, Right:%v", e.Left, e.Right)
}

func (e *UnaryNode) String() string {
	return fmt.Sprintf("Formula:%v", e.Formula)
}

func (e *QuantNode) String() string {
	return fmt.Sprintf("Variable:%v, %v", e.Variable, e.UnaryNode)
}

func (TrueNode) String() string {
	return "ExprTrue"
}

func (FalseNode) String() string {
	return "ExprFalse"
}

func (e *AndNode) String() string {
	return fmt.Sprintf("AndNode{%v}", e.BinaryNode)
}

func (e *OrNode) String() string {
	return fmt.Sprintf("OrNode{%v}", e.BinaryNode)
}

func (e *NotNode) String() string {
	return fmt.Sprintf("NotNode{%v}", e.UnaryNode)
}

func (e *ForallNode) String() string {
	return fmt.Sprintf("ForallNode{%v}", e.QuantNode)
}

func (e *ExistsNode) String() string {
	return fmt.Sprintf("ExistsNode[{%v}", e.QuantNode)
}

func (eff *WhenNode) String() string {
	return fmt.Sprintf("WhenNode{Condition:%v, }", eff.Condition,
		eff.UnaryNode)
}

func (eff *AssignNode) String() string {
	return fmt.Sprintf("AssignNode{Op:%v, Lval:%v, Rval:%v}",
		eff.Op, eff.Lval, eff.Rval)
}

var assignOpNames = map[AssignOp]string{
	OpAssign:    "OpAssign",
	OpScaleUp:   "OpScaleUp",
	OpScaleDown: "OpScaleDown",
	OpIncrease:  "OpIncrease",
	OpDecrease:  "OpDecrease",
}
<commit_msg>Automated merge with ssh://soong.cs.unh.edu/src/go/src/goplan<commit_after>package lifted

import (
	"fmt"
	"bytes"
)

func (d *Domain) String() string {
	var buf bytes.Buffer

	buf.WriteString("&Domain{")
	fmt.Fprintf(&buf, "Name:%s,\n", d.Name)
	fmt.Fprintf(&buf, "Requirements:%+v\n", d.Requirements)
	fmt.Fprintf(&buf, "Types:%v\n", d.Types)
	fmt.Fprintf(&buf, "Constants:%v\n", d.Constants)
	fmt.Fprintf(&buf, "Predicates:%+v\n", d.Predicates)

	buf.WriteString("Actions:[\n")
	for i, a := range d.Actions {
		buf.WriteString(a.String())
		if i < len(d.Actions)-1 {
			buf.WriteString("\n\n")
		}
	}
	buf.WriteString("],\n}")

	return buf.String()
}

func (p Predicate) String() string {
	return fmt.Sprintf("Predicate{Name:%v, Parameters:%v}",
		p.Name, p.Parameters)
}

func (a Action) String() string {
	var buf bytes.Buffer

	buf.WriteString("Action{")
	fmt.Fprintf(&buf, "Name:%s", a.Name)
	fmt.Fprintf(&buf, "\nParameters:%+v", a.Parameters)
	fmt.Fprintf(&buf, "\nPrecondition:%+v", a.Precondition)
	fmt.Fprintf(&buf, "\nEffect:%+v", a.Effect)
	buf.WriteString("}")

	return buf.String()
}

func (p *Problem) String() string {
	var buf bytes.Buffer

	buf.WriteString("Problem{")
	fmt.Fprintf(&buf, "Name:%s\n", p.Name)
	fmt.Fprintf(&buf, "Domain:%s\n", p.Domain)
	fmt.Fprintf(&buf, "Requirements:%v\n", p.Requirements)
	fmt.Fprintf(&buf, "Objects:%v\n", p.Objects)
	fmt.Fprintf(&buf, "Init:%+v\n", p.Init)
	fmt.Fprintf(&buf, "Goal:%+v\n", p.Goal)
	fmt.Fprintf(&buf, "Metric:%+v\n", p.Metric)
	buf.WriteByte('}')

	return buf.String()
}

func (k TermKind) String() string {
	switch k {
	case TermVariable:
		return "TermVariable"
	case TermConstant:
		return "TermConstant"
	}

	return fmt.Sprintf("%d", int(k))
}

func (t Term) String() string {
	return fmt.Sprintf("Term{Kind:%v, Name:%v}",
		t.Kind, t.Name)
}

func (n Name) String() string {
	if n.Num < 0 {
		return fmt.Sprintf("{%s}", n.Str)
	}
	return fmt.Sprintf("{%s, %d}", n.Str, n.Num)
}

func (lit *LiteralNode) String() string {
	return fmt.Sprintf("Literal{Positive:%t, Name:%v, Parameters:%v}",
		lit.Positive, lit.Name, lit.Parameters)
}

func (e *BinaryNode) String() string {
	return fmt.Sprintf("Left:%v, Right:%v", e.Left, e.Right)
}

func (e *UnaryNode) String() string {
	return fmt.Sprintf("Formula:%v", e.Formula)
}

func (e *QuantNode) String() string {
	return fmt.Sprintf("Variable:%v, %v", e.Variable, e.UnaryNode)
}

func (TrueNode) String() string {
	return "ExprTrue"
}

func (FalseNode) String() string {
	return "ExprFalse"
}

func (e *AndNode) String() string {
	return fmt.Sprintf("AndNode{%v}", e.BinaryNode)
}

func (e *OrNode) String() string {
	return fmt.Sprintf("OrNode{%v}", e.BinaryNode)
}

func (e *NotNode) String() string {
	return fmt.Sprintf("NotNode{%v}", e.UnaryNode)
}

func (e *ForallNode) String() string {
	return fmt.Sprintf("ForallNode{%v}", e.QuantNode)
}

func (e *ExistsNode) String() string {
	return fmt.Sprintf("ExistsNode[{%v}", e.QuantNode)
}

func (eff *WhenNode) String() string {
	return fmt.Sprintf("WhenNode{Condition:%v, }", eff.Condition,
		eff.UnaryNode)
}

func (eff *AssignNode) String() string {
	return fmt.Sprintf("AssignNode{Op:%v, Lval:%v, Rval:%v}",
		eff.Op, eff.Lval, eff.Rval)
}

var assignOpNames = map[AssignOp]string{
	OpAssign:    "OpAssign",
	OpScaleUp:   "OpScaleUp",
	OpScaleDown: "OpScaleDown",
	OpIncrease:  "OpIncrease",
	OpDecrease:  "OpDecrease",
}
<|endoftext|>


<commit_before>/* Copyright 2016 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gonids

import (
	"testing"
)

func TestShouldBeHTTP(t *testing.T) {
	for _, tt := range []struct {
		name  string
		input *Rule
		want  bool
	}{
		{
			name: "already http",
			input: &Rule{
				Protocol: "http",
			},
			want: false,
		},
		{
			name: "content option change",
			input: &Rule{
				Protocol: "tcp",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
				},
			},
			want: true,
		},
		{
			name: "sticky buffer change",
			input: &Rule{
				Protocol: "tcp",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						DataPosition: httpProtocol,
						Pattern:      []byte("AA"),
					},
				},
			},
			want: true,
		},
	} {
		got := tt.input.OptimizeHTTP()
		// Expected modification.
		if got != tt.want {
			t.Fatalf("%s: got %v; want %v", tt.name, got, tt.want)
		}
	}
}

func TestExpensivePCRE(t *testing.T) {
	for _, tt := range []struct {
		name  string
		input *Rule
		want  bool
	}{
		{
			name: "No PCRE",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AAAAAAAAAA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
				},
			},
			want: false,
		},
		{
			name: "No Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Short Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Only Common Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("POST"),
						Options: []*ContentOption{
							{"http_method", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Long Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("ReallyLongThing"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: false,
		},
		{
			name: "Banned complex content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("\r\nUser-Agent: "),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
	} {
		got := tt.input.ExpensivePCRE()
		// Expected modification.
		if got != tt.want {
			t.Fatalf("%s: got %v; want %v", tt.name, got, tt.want)
		}
	}
}
<commit_msg>Add another test.<commit_after>/* Copyright 2016 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gonids

import (
	"testing"
)

func TestShouldBeHTTP(t *testing.T) {
	for _, tt := range []struct {
		name  string
		input *Rule
		want  bool
	}{
		{
			name: "already http",
			input: &Rule{
				Protocol: "http",
			},
			want: false,
		},
		{
			name: "content option change",
			input: &Rule{
				Protocol: "tcp",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
				},
			},
			want: true,
		},
		{
			name: "sticky buffer change",
			input: &Rule{
				Protocol: "tcp",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						DataPosition: httpProtocol,
						Pattern:      []byte("AA"),
					},
				},
			},
			want: true,
		},
	} {
		got := tt.input.OptimizeHTTP()
		// Expected modification.
		if got != tt.want {
			t.Fatalf("%s: got %v; want %v", tt.name, got, tt.want)
		}
	}
}

func TestExpensivePCRE(t *testing.T) {
	for _, tt := range []struct {
		name  string
		input *Rule
		want  bool
	}{
		{
			name: "No PCRE",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AAAAAAAAAA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
				},
			},
			want: false,
		},
		{
			name: "No Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Short Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("AA"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Only Common Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("POST"),
						Options: []*ContentOption{
							{"http_method", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Long Content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("ReallyLongThing"),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: false,
		},
		{
			name: "Banned complex content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("\r\nUser-Agent: "),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: true,
		},
		{
			name: "Banned complex content, with long content",
			input: &Rule{
				Protocol: "http",
				Source: Network{
					Nets:  []string{"$HOME_NET"},
					Ports: []string{"any"},
				},
				Destination: Network{
					Nets:  []string{"$EXTERNAL_NET"},
					Ports: []string{"$HTTP_PORTS"},
				},
				Matchers: []orderedMatcher{
					&Content{
						Pattern: []byte("\r\nUser-Agent: "),
						Options: []*ContentOption{
							{"http_header", ""},
						},
					},
					&Content{
						Pattern: []byte("SuperLongUniqueAwesome"),
					},
					&PCRE{
						Pattern: []byte("f.*bar"),
					},
				},
			},
			want: false,
		},
	} {
		got := tt.input.ExpensivePCRE()
		// Expected modification.
		if got != tt.want {
			t.Fatalf("%s: got %v; want %v", tt.name, got, tt.want)
		}
	}
}
<|endoftext|>


<commit_before>package pe

import (
	"bufio"
	"bytes"
	"debug/pe"
	"encoding/binary"
	"fmt"
	"github.com/williballenthin/Lancelot/workspace"
	"unicode/utf16"
)

var PAGE_SIZE uint64 = 0x1000

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func roundUp(i uint64, base uint64) uint64 {
	if i%base == 0x0 {
		return i
	} else {
		return i + base - (i % base)
	}
}

func roundUpToPage(i uint64) uint64 {
	return roundUp(i, PAGE_SIZE)
}

type PELoader struct {
	name string
	file *pe.File
}

func New(name string, file *pe.File) (*PELoader, error) {
	// TODO: detect bitness
	return &PELoader{name: name, file: file}, nil
}

type ImageImportDirectory struct {
	rvaOriginalThunkTable uint32
	TimeDateStamp         uint32
	ForwarderChain        uint32
	rvaModuleName         uint32
	rvaThunkTable         uint32
}

func (loader *PELoader) loadPESection(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	section *pe.Section) error {

	h := section.SectionHeader

	fmt.Printf("section: %s\n", section.SectionHeader.Name)
	fmt.Printf("  virtual address: 0x%x\n", section.SectionHeader.VirtualAddress)
	fmt.Printf("  virtual size: 0x%x\n", section.SectionHeader.VirtualSize)
	fmt.Printf("  file offset: 0x%x\n", section.SectionHeader.Offset)
	fmt.Printf("  file size: 0x%x\n", section.SectionHeader.Size)

	rvaSecStart := workspace.RVA(h.VirtualAddress)
	secStart := mod.VA(rvaSecStart)
	secLength := roundUpToPage(uint64(h.VirtualSize))
	e := ws.MemMap(secStart, secLength, fmt.Sprintf("%s/%s", mod.Name, section.SectionHeader.Name))
	check(e)

	d, e := section.Data()
	check(e)

	e = mod.MemWrite(ws, rvaSecStart, d)
	check(e)

	// TODO: apply permissions

	return nil
}

type ImageImportByName struct {
	Hint uint16
	Name string
}

type Flag64 uint64
type Flags64 uint64

var FLAG_IMPORT_BY_ORDINAL Flag64 = 1 << 31

func (i Flags64) isSet(j Flag64) bool {
	return uint64(i)&uint64(j) > 0
}

func (loader *PELoader) resolveThunkTable(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	moduleName string,
	rvaTable workspace.RVA) error {

	var offset workspace.RVA = rvaTable
	for {
		rvaImport, e := mod.MemReadRva(ws, offset)
		check(e)

		if rvaImport == 0x0 {
			break
		}

		if Flags64(rvaImport).isSet(FLAG_IMPORT_BY_ORDINAL) {
			fmt.Printf("  import by ordinal: %03x\n", uint64(rvaImport)&uint64(0x7FFFFFFF))
			// TODO: replace thunk with handler
			// notes:
			//    32: PUSH 0xAABBCCDD --> 68 DD CC BB AA
			//        JMP  0xAABBCCDD --> E9 D9 CC BB AA  ugh, relative jump. do a push/ret instead.
			//        RET             --> C3
			//
			mod.Imports[offset] = workspace.LinkedSymbol{
				ModuleName: moduleName,
				SymbolName: fmt.Sprintf("ordinal-%x", uint64(rvaImport)&uint64(0x7FFFFFFF)),
			}
		} else {
			d, e := mod.MemRead(ws, rvaImport, 0x100)
			check(e)

			p := bytes.NewBuffer(d)
			var importByName ImageImportByName
			binary.Read(p, binary.LittleEndian, &importByName.Hint)

			importByName.Name, e = readAscii(d[2:])
			check(e)

			fmt.Printf("  import by name: %s@0x%x\n", importByName.Name, uint64(rvaImport))
			// TODO: replace thunk with handler
			mod.Imports[offset] = workspace.LinkedSymbol{
				ModuleName: moduleName,
				SymbolName: importByName.Name,
			}
		}

		offset += 4
	}
	return nil
}

func (loader *PELoader) resolveImports(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	dataDirectory [16]pe.DataDirectory) error {

	// since we always map at ImageBase, we don't need to apply (32bit) relocs
	// TODO: check 64bit reloc types

	importDirectory := dataDirectory[1]
	importRva := workspace.RVA(importDirectory.VirtualAddress)
	importSize := importDirectory.Size

	fmt.Printf("import rva: 0x%x\n", importRva)
	fmt.Printf("import size: 0x%x\n", importSize)

	d, e := mod.MemRead(ws, workspace.RVA(importDirectory.VirtualAddress), uint64(importDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	for {
		var dir ImageImportDirectory
		binary.Read(p, binary.LittleEndian, &dir.rvaOriginalThunkTable)
		fmt.Printf("rva import lookup table: 0x%x\n", dir.rvaOriginalThunkTable)
		if dir.rvaOriginalThunkTable == 0 {
			break
		}
		binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
		fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

		binary.Read(p, binary.LittleEndian, &dir.ForwarderChain)
		fmt.Printf("forwarder chain: 0x%x\n", dir.ForwarderChain)

		binary.Read(p, binary.LittleEndian, &dir.rvaModuleName)

		moduleNameBuf, e := mod.MemRead(ws, workspace.RVA(dir.rvaModuleName), 0x100)
		check(e)
		moduleName, e := readAscii(moduleNameBuf)
		check(e)

		fmt.Printf("module name: %s\n", string(moduleName))

		binary.Read(p, binary.LittleEndian, &dir.rvaThunkTable)
		loader.resolveThunkTable(ws, mod, moduleName, workspace.RVA(dir.rvaThunkTable))
	}

	return nil
}

type ImageExportDirectory struct {
	Characteristics          uint32
	TimeDateStamp            uint32
	MajorVersion             uint16
	MinorVersion             uint16
	rvaName                  uint32
	Base                     uint32
	NumberOfFunctions        uint32
	NumberOfNames            uint32
	rvaAddressOfFunctions    uint32
	rvaAddressOfNames        uint32
	rvaAddressOfNameOrdinals uint32
}

func (loader *PELoader) resolveExports(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	dataDirectory [16]pe.DataDirectory) error {
	exportDirectory := dataDirectory[0]
	exportRva := workspace.RVA(exportDirectory.VirtualAddress)
	exportSize := exportDirectory.Size

	fmt.Printf("export rva: 0x%x\n", exportRva)
	fmt.Printf("export size: 0x%x\n", exportSize)

	d, e := mod.MemRead(ws, workspace.RVA(exportDirectory.VirtualAddress), uint64(exportDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	var dir ImageExportDirectory

	binary.Read(p, binary.LittleEndian, &dir.Characteristics)
	binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
	binary.Read(p, binary.LittleEndian, &dir.MajorVersion)
	binary.Read(p, binary.LittleEndian, &dir.MinorVersion)
	binary.Read(p, binary.LittleEndian, &dir.rvaName)
	binary.Read(p, binary.LittleEndian, &dir.Base)
	binary.Read(p, binary.LittleEndian, &dir.NumberOfFunctions)
	binary.Read(p, binary.LittleEndian, &dir.NumberOfNames)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfFunctions)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfNames)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfNameOrdinals)

	if dir.rvaAddressOfFunctions == 0 {
		panic("address of functions is NULL")
	}

	exportModuleNameBuf, e := mod.MemRead(ws, workspace.RVA(dir.rvaName), 0x100)
	check(e)
	exportModuleName, e := readAscii(exportModuleNameBuf)
	fmt.Printf("export name: %ss\n", string(exportModuleName))
	check(e)

	fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

	// resolve exports by ordinals first
	for i := uint32(0); i < dir.NumberOfFunctions; i++ {
		ordinal := uint16(i + dir.Base)
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaFn, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfFunctions+4*i))
		check(e)
		mod.ExportsByOrdinal[ordinal] = rvaFn
		fmt.Printf(" export: (ordinal) %x: 0x%x\n", ordinal, rvaFn)
	}

	// resolve exports by name
	for i := uint32(0); i < dir.NumberOfNames; i++ {
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaName, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfNames+4*i))
		check(e)
		// sizeof(ordinal) is always 2 bytes
		nameOrdinal, e := mod.MemReadShort(ws, workspace.RVA(dir.rvaAddressOfNameOrdinals+2*i))
		check(e)
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaFn, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfFunctions+4*uint32(nameOrdinal)))
		check(e)

		nameBuf, e := mod.MemRead(ws, workspace.RVA(rvaName), 0x100)
		check(e)
		name, e := readAscii(nameBuf)
		check(e)

		mod.ExportsByName[name] = rvaFn
		fmt.Printf(" export: %s: 0x%x\n", name, rvaFn)
	}

	return nil
}

func (loader *PELoader) Load(ws *workspace.Workspace) (*workspace.LoadedModule, error) {
	var imageBase workspace.VA
	var addressOfEntryPoint workspace.RVA
	var dataDirectory [16]pe.DataDirectory

	if optionalHeader, ok := loader.file.OptionalHeader.(*pe.OptionalHeader32); ok {
		imageBase = workspace.VA(optionalHeader.ImageBase)
		addressOfEntryPoint = workspace.RVA(optionalHeader.AddressOfEntryPoint)
		dataDirectory = optionalHeader.DataDirectory
	} else {
		return nil, workspace.InvalidModeError
	}

	mod := &workspace.LoadedModule{
		Name:             loader.name,
		BaseAddress:      imageBase,
		EntryPoint:       addressOfEntryPoint.VA(imageBase),
		Imports:          map[workspace.RVA]workspace.LinkedSymbol{},
		ExportsByName:    map[string]workspace.RVA{},
		ExportsByOrdinal: map[uint16]workspace.RVA{},
	}

	for _, section := range loader.file.Sections {
		e := loader.loadPESection(ws, mod, section)
		check(e)
	}

	e := loader.resolveImports(ws, mod, dataDirectory)
	check(e)

	e = loader.resolveExports(ws, mod, dataDirectory)
	check(e)

	e = ws.AddLoadedModule(mod)
	check(e)

	return mod, nil
}

func readAscii(buf []byte) (string, error) {
	br := bufio.NewReader(bytes.NewReader(buf))
	bytez, e := br.ReadBytes(byte(0x00))
	check(e)
	return string(bytez), nil
}

func readUtf16le(buf []byte) ([]rune, error) {
	start := 0
	end := 0

	for i := 0; i < len(buf)/2; i++ {
		if !(buf[i*2] == 0 && buf[i*2+1] == 0) {
			end = i * 2
			break
		}
	}

	numChars := (end - start) / 2
	d := make([]uint16, numChars)
	for i := 0; i < numChars; i++ {
		d[i] = uint16(buf[i*2]) | (uint16(buf[i*2+1]) << 8)
	}
	return utf16.Decode(d), nil
}
<commit_msg>pe: trim trailing NULL bytes from ASCII strings<commit_after>package pe

import (
	"bufio"
	"bytes"
	"debug/pe"
	"encoding/binary"
	"fmt"
	"github.com/williballenthin/Lancelot/workspace"
	"unicode/utf16"
)

var PAGE_SIZE uint64 = 0x1000

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func roundUp(i uint64, base uint64) uint64 {
	if i%base == 0x0 {
		return i
	} else {
		return i + base - (i % base)
	}
}

func roundUpToPage(i uint64) uint64 {
	return roundUp(i, PAGE_SIZE)
}

type PELoader struct {
	name string
	file *pe.File
}

func New(name string, file *pe.File) (*PELoader, error) {
	// TODO: detect bitness
	return &PELoader{name: name, file: file}, nil
}

type ImageImportDirectory struct {
	rvaOriginalThunkTable uint32
	TimeDateStamp         uint32
	ForwarderChain        uint32
	rvaModuleName         uint32
	rvaThunkTable         uint32
}

func (loader *PELoader) loadPESection(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	section *pe.Section) error {

	h := section.SectionHeader

	fmt.Printf("section: %s\n", section.SectionHeader.Name)
	fmt.Printf("  virtual address: 0x%x\n", section.SectionHeader.VirtualAddress)
	fmt.Printf("  virtual size: 0x%x\n", section.SectionHeader.VirtualSize)
	fmt.Printf("  file offset: 0x%x\n", section.SectionHeader.Offset)
	fmt.Printf("  file size: 0x%x\n", section.SectionHeader.Size)

	rvaSecStart := workspace.RVA(h.VirtualAddress)
	secStart := mod.VA(rvaSecStart)
	secLength := roundUpToPage(uint64(h.VirtualSize))
	e := ws.MemMap(secStart, secLength, fmt.Sprintf("%s/%s", mod.Name, section.SectionHeader.Name))
	check(e)

	d, e := section.Data()
	check(e)

	e = mod.MemWrite(ws, rvaSecStart, d)
	check(e)

	// TODO: apply permissions

	return nil
}

type ImageImportByName struct {
	Hint uint16
	Name string
}

type Flag64 uint64
type Flags64 uint64

var FLAG_IMPORT_BY_ORDINAL Flag64 = 1 << 31

func (i Flags64) isSet(j Flag64) bool {
	return uint64(i)&uint64(j) > 0
}

func (loader *PELoader) resolveThunkTable(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	moduleName string,
	rvaTable workspace.RVA) error {

	var offset workspace.RVA = rvaTable
	for {
		rvaImport, e := mod.MemReadRva(ws, offset)
		check(e)

		if rvaImport == 0x0 {
			break
		}

		if Flags64(rvaImport).isSet(FLAG_IMPORT_BY_ORDINAL) {
			fmt.Printf("  import by ordinal: %03x\n", uint64(rvaImport)&uint64(0x7FFFFFFF))
			// TODO: replace thunk with handler
			// notes:
			//    32: PUSH 0xAABBCCDD --> 68 DD CC BB AA
			//        JMP  0xAABBCCDD --> E9 D9 CC BB AA  ugh, relative jump. do a push/ret instead.
			//        RET             --> C3
			//
			mod.Imports[offset] = workspace.LinkedSymbol{
				ModuleName: moduleName,
				SymbolName: fmt.Sprintf("ordinal-%x", uint64(rvaImport)&uint64(0x7FFFFFFF)),
			}
		} else {
			d, e := mod.MemRead(ws, rvaImport, 0x100)
			check(e)

			p := bytes.NewBuffer(d)
			var importByName ImageImportByName
			binary.Read(p, binary.LittleEndian, &importByName.Hint)

			importByName.Name, e = readAscii(d[2:])
			check(e)

			fmt.Printf("  import by name: %s@0x%x\n", importByName.Name, uint64(rvaImport))
			// TODO: replace thunk with handler
			mod.Imports[offset] = workspace.LinkedSymbol{
				ModuleName: moduleName,
				SymbolName: importByName.Name,
			}
		}

		offset += 4
	}
	return nil
}

func (loader *PELoader) resolveImports(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	dataDirectory [16]pe.DataDirectory) error {

	// since we always map at ImageBase, we don't need to apply (32bit) relocs
	// TODO: check 64bit reloc types

	importDirectory := dataDirectory[1]
	importRva := workspace.RVA(importDirectory.VirtualAddress)
	importSize := importDirectory.Size

	fmt.Printf("import rva: 0x%x\n", importRva)
	fmt.Printf("import size: 0x%x\n", importSize)

	d, e := mod.MemRead(ws, workspace.RVA(importDirectory.VirtualAddress), uint64(importDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	for {
		var dir ImageImportDirectory
		binary.Read(p, binary.LittleEndian, &dir.rvaOriginalThunkTable)
		fmt.Printf("rva import lookup table: 0x%x\n", dir.rvaOriginalThunkTable)
		if dir.rvaOriginalThunkTable == 0 {
			break
		}
		binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
		fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

		binary.Read(p, binary.LittleEndian, &dir.ForwarderChain)
		fmt.Printf("forwarder chain: 0x%x\n", dir.ForwarderChain)

		binary.Read(p, binary.LittleEndian, &dir.rvaModuleName)

		moduleNameBuf, e := mod.MemRead(ws, workspace.RVA(dir.rvaModuleName), 0x100)
		check(e)
		moduleName, e := readAscii(moduleNameBuf)
		check(e)

		fmt.Printf("module name: %s\n", string(moduleName))

		binary.Read(p, binary.LittleEndian, &dir.rvaThunkTable)
		loader.resolveThunkTable(ws, mod, moduleName, workspace.RVA(dir.rvaThunkTable))
	}

	return nil
}

type ImageExportDirectory struct {
	Characteristics          uint32
	TimeDateStamp            uint32
	MajorVersion             uint16
	MinorVersion             uint16
	rvaName                  uint32
	Base                     uint32
	NumberOfFunctions        uint32
	NumberOfNames            uint32
	rvaAddressOfFunctions    uint32
	rvaAddressOfNames        uint32
	rvaAddressOfNameOrdinals uint32
}

func (loader *PELoader) resolveExports(
	ws *workspace.Workspace,
	mod *workspace.LoadedModule,
	dataDirectory [16]pe.DataDirectory) error {
	exportDirectory := dataDirectory[0]
	exportRva := workspace.RVA(exportDirectory.VirtualAddress)
	exportSize := exportDirectory.Size

	fmt.Printf("export rva: 0x%x\n", exportRva)
	fmt.Printf("export size: 0x%x\n", exportSize)

	d, e := mod.MemRead(ws, workspace.RVA(exportDirectory.VirtualAddress), uint64(exportDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	var dir ImageExportDirectory

	binary.Read(p, binary.LittleEndian, &dir.Characteristics)
	binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
	binary.Read(p, binary.LittleEndian, &dir.MajorVersion)
	binary.Read(p, binary.LittleEndian, &dir.MinorVersion)
	binary.Read(p, binary.LittleEndian, &dir.rvaName)
	binary.Read(p, binary.LittleEndian, &dir.Base)
	binary.Read(p, binary.LittleEndian, &dir.NumberOfFunctions)
	binary.Read(p, binary.LittleEndian, &dir.NumberOfNames)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfFunctions)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfNames)
	binary.Read(p, binary.LittleEndian, &dir.rvaAddressOfNameOrdinals)

	if dir.rvaAddressOfFunctions == 0 {
		panic("address of functions is NULL")
	}

	exportModuleNameBuf, e := mod.MemRead(ws, workspace.RVA(dir.rvaName), 0x100)
	check(e)
	exportModuleName, e := readAscii(exportModuleNameBuf)
	fmt.Printf("export name: %ss\n", string(exportModuleName))
	check(e)

	fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

	// resolve exports by ordinals first
	for i := uint32(0); i < dir.NumberOfFunctions; i++ {
		ordinal := uint16(i + dir.Base)
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaFn, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfFunctions+4*i))
		check(e)
		mod.ExportsByOrdinal[ordinal] = rvaFn
		fmt.Printf(" export: (ordinal) %x: 0x%x\n", ordinal, rvaFn)
	}

	// resolve exports by name
	for i := uint32(0); i < dir.NumberOfNames; i++ {
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaName, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfNames+4*i))
		check(e)
		// sizeof(ordinal) is always 2 bytes
		nameOrdinal, e := mod.MemReadShort(ws, workspace.RVA(dir.rvaAddressOfNameOrdinals+2*i))
		check(e)
		// sizeof(RVA) is always 4 bytes, even on x64
		rvaFn, e := mod.MemReadRva(ws, workspace.RVA(dir.rvaAddressOfFunctions+4*uint32(nameOrdinal)))
		check(e)

		nameBuf, e := mod.MemRead(ws, workspace.RVA(rvaName), 0x100)
		check(e)
		name, e := readAscii(nameBuf)
		check(e)

		mod.ExportsByName[name] = rvaFn
		fmt.Printf(" export: %s: 0x%x\n", name, rvaFn)
	}

	return nil
}

func (loader *PELoader) Load(ws *workspace.Workspace) (*workspace.LoadedModule, error) {
	var imageBase workspace.VA
	var addressOfEntryPoint workspace.RVA
	var dataDirectory [16]pe.DataDirectory

	if optionalHeader, ok := loader.file.OptionalHeader.(*pe.OptionalHeader32); ok {
		imageBase = workspace.VA(optionalHeader.ImageBase)
		addressOfEntryPoint = workspace.RVA(optionalHeader.AddressOfEntryPoint)
		dataDirectory = optionalHeader.DataDirectory
	} else {
		return nil, workspace.InvalidModeError
	}

	mod := &workspace.LoadedModule{
		Name:             loader.name,
		BaseAddress:      imageBase,
		EntryPoint:       addressOfEntryPoint.VA(imageBase),
		Imports:          map[workspace.RVA]workspace.LinkedSymbol{},
		ExportsByName:    map[string]workspace.RVA{},
		ExportsByOrdinal: map[uint16]workspace.RVA{},
	}

	for _, section := range loader.file.Sections {
		e := loader.loadPESection(ws, mod, section)
		check(e)
	}

	e := loader.resolveImports(ws, mod, dataDirectory)
	check(e)

	e = loader.resolveExports(ws, mod, dataDirectory)
	check(e)

	e = ws.AddLoadedModule(mod)
	check(e)

	return mod, nil
}

func readAscii(buf []byte) (string, error) {
	br := bufio.NewReader(bytes.NewReader(buf))
	bytez, e := br.ReadBytes(byte(0x00))
	check(e)
	bytez = bytes.TrimSuffix(bytez, []byte{0x00})
	return string(bytez), nil
}

func readUtf16le(buf []byte) ([]rune, error) {
	start := 0
	end := 0

	for i := 0; i < len(buf)/2; i++ {
		if !(buf[i*2] == 0 && buf[i*2+1] == 0) {
			end = i * 2
			break
		}
	}

	numChars := (end - start) / 2
	d := make([]uint16, numChars)
	for i := 0; i < numChars; i++ {
		d[i] = uint16(buf[i*2]) | (uint16(buf[i*2+1]) << 8)
	}
	return utf16.Decode(d), nil
}
<|endoftext|>