<commit_before>package main

import (
	"bufio"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"golang.org/x/sync/errgroup"
	"os"
	"reflect"
	"regexp"
	"strconv"
)

// reconstructOptions represents all options for reconstruction
type reconstructOptions struct {
	BpP    int
	LimitX int
	LimitY int
	Scale  int
	Dtg    string
	Source string
	Filter string
}

// svgOptions represents various options for reconstruction
type svgOptions struct {
	regex             string
	reconstructOption string
}

func createPacket(ch chan<- []byte, packet []int, bpP int) error {
	var buf []byte
	var tmp int
	switch bpP {
	case 24:
		for _, i := range packet {
			buf = append(buf, byte(i))
		}
	case 3, 6, 9, 12, 15, 18, 21:
		var slice []int
		for i := 0; i < len(packet); i = i + 1 {
			if i%(bpP*8) == 0 && i != 0 {
				bytes := createBytes(slice, bpP/3)
				buf = append(buf, bytes...)
				slice = slice[:0]
			}
			slice = append(slice, packet[i])
		}
		bytes := createBytes(slice, bpP/3)
		buf = append(buf, bytes...)
	case 1:
		var j int
		for i := 0; i < len(packet); i = i + 3 {
			if j%8 == 0 && j != 0 {
				buf = append(buf, byte(tmp))
				tmp = 0
			}
			if packet[i] != 0 {
				tmp = tmp | (1 << uint8(7-j%8))
			}
			j = j + 1
		}
		if tmp != 0 {
			buf = append(buf, byte(tmp))
		}
	default:
		return fmt.Errorf("This format is not supported so far")
	}

	ch <- buf

	return nil
}

func checkVersion(parse *[]svgOptions, version string) (string, error) {

	scale := svgOptions{regex: "\\s+Scale=(\\d+)$", reconstructOption: "Scale"}
	*parse = append(*parse, scale)
	bpP := svgOptions{regex: "\\s+BitsPerPixel=(\\d+)$", reconstructOption: "BpP"}
	*parse = append(*parse, bpP)

	switch version {
	case "0.0.3":
		dtg := svgOptions{regex: "\\s+DTG=\"([0-9. :a-zA-Z]+)\"", reconstructOption: "Dtg"}
		*parse = append(*parse, dtg)
		source := svgOptions{regex: "\\s+Source=\"(\\w+)\"", reconstructOption: "Source"}
		*parse = append(*parse, source)
		filter := svgOptions{regex: "\\s+Filter=\"(\\w+)\"", reconstructOption: "Filter"}
		*parse = append(*parse, filter)
	case "0.0.4":
		lGate := svgOptions{regex: "\\s+LogicGate=\"([a-zA-Z]+)\"", reconstructOption: "LogicGate"}
		*parse = append(*parse, lGate)
		lValue := svgOptions{regex: "\\s+LogicValue=(0x[0-9A-F]{2})", reconstructOption: "LogicValue"}
		*parse = append(*parse, lValue)
		return "", fmt.Errorf("Can't decode version %s at the moment.")
	default:
		return "", fmt.Errorf("Unrecognized version: %s", version)
	}
	return version, nil
}

func checkHeader(svg *bufio.Scanner) (reconstructOptions, error) {
	var options reconstructOptions
	var variant string
	var header bool = false
	var parseOptions []svgOptions
	var optionIndex int

	limits, err := regexp.Compile("^<svg width=\"(\\d+)\" height=\"(\\d+)\">$")
	if err != nil {
		return options, err
	}
	headerStart, err := regexp.Compile("^<!--$")
	if err != nil {
		return options, err
	}
	headerEnd, err := regexp.Compile("^-->$")
	if err != nil {
		return options, err
	}
	version, err := regexp.Compile("\\s+goNetViz \"([0-9.]+)\"$")
	if err != nil {
		return options, err
	}

	for svg.Scan() {
		line := svg.Text()
		switch {
		case options.LimitX == 0 && options.LimitY == 0 && header == false:
			matches := limits.FindStringSubmatch(line)
			if len(matches) == 3 {
				options.LimitX, _ = strconv.Atoi(matches[1])
				options.LimitY, _ = strconv.Atoi(matches[2])
			}
		case header == false:
			if headerStart.MatchString(line) {
				header = true
			}
		case len(variant) == 0:
			matches := version.FindStringSubmatch(line)
			if len(matches) == 2 {
				variant, err = checkVersion(&parseOptions, matches[1])
				if err != nil {
					return options, err
				}
			}
		default:
			if optionIndex > len(parseOptions) {
				return options, fmt.Errorf("Option index is out of range")
			}
			regex, err := regexp.Compile(parseOptions[optionIndex].regex)
			if err != nil {
				return options, err
			}
			matches := regex.FindStringSubmatch(line)
			if len(matches) == 2 {
				option := reflect.ValueOf(&options).Elem().FieldByName(parseOptions[optionIndex].reconstructOption)

				switch option.Kind() {
				case reflect.Int:
					new, _ := strconv.Atoi(matches[1])
					option.SetInt(int64(new))
				case reflect.String:
					option.SetString(matches[1])
				default:
					return options, fmt.Errorf("Unhandeld option type")
				}
				optionIndex += 1
			} else {
				if headerEnd.MatchString(line) {
					return options, nil
				}
			}
		}
	}
	return options, fmt.Errorf("No end of header found")
}

func extractInformation(g *errgroup.Group, ch chan []byte, cfg configs) error {
	inputfile, err := os.Open(cfg.file)
	if err != nil {
		return fmt.Errorf("Could not open file %s: %s\n", cfg.file, err.Error())
	}
	defer inputfile.Close()
	svg := bufio.NewScanner(inputfile)
	var yLast int
	var packet []int
	defer close(ch)

	opt, err := checkHeader(svg)
	if err != nil {
		return err
	}

	pixel, err := regexp.Compile("^<rect x=\"(\\d+)\" y=\"(\\d+)\" width=\"\\d+\" height=\"\\d+\" style=\"fill:rgb\\((\\d+),(\\d+),(\\d+)\\)\" />$")
	if err != nil {
		return err
	}
	svgEnd, err := regexp.Compile("</svg>")
	if err != nil {
		return err
	}

	for svg.Scan() {
		line := svg.Text()
		matches := pixel.FindStringSubmatch(line)
		if len(matches) == 6 {
			pixelX, _ := strconv.Atoi(matches[1])
			pixelY, _ := strconv.Atoi(matches[2])
			if pixelY != yLast {
				yLast = pixelY
				if err := createPacket(ch, packet, opt.BpP); err != nil {
					return err
				}
				packet = packet[:0]
			}
			if pixelX >= opt.LimitX {
				return fmt.Errorf("x-coordinate (%d) is bigger than the limit (%d)\n", pixelX, opt.LimitX)
			}
			r, _ := strconv.Atoi(matches[3])
			g, _ := strconv.Atoi(matches[4])
			b, _ := strconv.Atoi(matches[5])
			packet = append(packet, r, g, b)
		} else if svgEnd.MatchString(line) {
			if len(packet) != 0 {
				return createPacket(ch, packet, opt.BpP)
			}
		}
	}
	return nil
}

func createPcap(g *errgroup.Group, ch chan []byte, cfg configs) error {
	filename := cfg.prefix
	filename += ".pcap"
	output, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("Could not create file %s: %s\n", filename, err.Error())
	}
	defer output.Close()
	w := pcapgo.NewWriter(output)
	w.WriteFileHeader(65536, layers.LinkTypeEthernet)

	for i, ok := <-ch; ok; i, ok = <-ch {
		w.WritePacket(gopacket.CaptureInfo{CaptureLength: len(i), Length: len(i), InterfaceIndex: 0}, i)
	}

	return nil
}

func reconstruct(g *errgroup.Group, cfg configs) error {
	ch := make(chan []byte)

	go extractInformation(g, ch, cfg)

	g.Go(func() error {
		return createPcap(g, ch, cfg)
	})

	return g.Wait()
}
<commit_msg>Fix error<commit_after>package main

import (
	"bufio"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"golang.org/x/sync/errgroup"
	"os"
	"reflect"
	"regexp"
	"strconv"
)

// reconstructOptions represents all options for reconstruction
type reconstructOptions struct {
	BpP    int
	LimitX int
	LimitY int
	Scale  int
	Dtg    string
	Source string
	Filter string
}

// svgOptions represents various options for reconstruction
type svgOptions struct {
	regex             string
	reconstructOption string
}

func createPacket(ch chan<- []byte, packet []int, bpP int) error {
	var buf []byte
	var tmp int
	switch bpP {
	case 24:
		for _, i := range packet {
			buf = append(buf, byte(i))
		}
	case 3, 6, 9, 12, 15, 18, 21:
		var slice []int
		for i := 0; i < len(packet); i = i + 1 {
			if i%(bpP*8) == 0 && i != 0 {
				bytes := createBytes(slice, bpP/3)
				buf = append(buf, bytes...)
				slice = slice[:0]
			}
			slice = append(slice, packet[i])
		}
		bytes := createBytes(slice, bpP/3)
		buf = append(buf, bytes...)
	case 1:
		var j int
		for i := 0; i < len(packet); i = i + 3 {
			if j%8 == 0 && j != 0 {
				buf = append(buf, byte(tmp))
				tmp = 0
			}
			if packet[i] != 0 {
				tmp = tmp | (1 << uint8(7-j%8))
			}
			j = j + 1
		}
		if tmp != 0 {
			buf = append(buf, byte(tmp))
		}
	default:
		return fmt.Errorf("This format is not supported so far")
	}

	ch <- buf

	return nil
}

func checkVersion(parse *[]svgOptions, version string) (string, error) {

	scale := svgOptions{regex: "\\s+Scale=(\\d+)$", reconstructOption: "Scale"}
	*parse = append(*parse, scale)
	bpP := svgOptions{regex: "\\s+BitsPerPixel=(\\d+)$", reconstructOption: "BpP"}
	*parse = append(*parse, bpP)

	switch version {
	case "0.0.3":
		dtg := svgOptions{regex: "\\s+DTG=\"([0-9. :a-zA-Z]+)\"", reconstructOption: "Dtg"}
		*parse = append(*parse, dtg)
		source := svgOptions{regex: "\\s+Source=\"(\\w+)\"", reconstructOption: "Source"}
		*parse = append(*parse, source)
		filter := svgOptions{regex: "\\s+Filter=\"(\\w+)\"", reconstructOption: "Filter"}
		*parse = append(*parse, filter)
	case "0.0.4":
		lGate := svgOptions{regex: "\\s+LogicGate=\"([a-zA-Z]+)\"", reconstructOption: "LogicGate"}
		*parse = append(*parse, lGate)
		lValue := svgOptions{regex: "\\s+LogicValue=(0x[0-9A-F]{2})", reconstructOption: "LogicValue"}
		*parse = append(*parse, lValue)
		return "", fmt.Errorf("Can't decode version 0.0.4 at the moment.")
	default:
		return "", fmt.Errorf("Unrecognized version: %s", version)
	}
	return version, nil
}

func checkHeader(svg *bufio.Scanner) (reconstructOptions, error) {
	var options reconstructOptions
	var variant string
	var header bool = false
	var parseOptions []svgOptions
	var optionIndex int

	limits, err := regexp.Compile("^<svg width=\"(\\d+)\" height=\"(\\d+)\">$")
	if err != nil {
		return options, err
	}
	headerStart, err := regexp.Compile("^<!--$")
	if err != nil {
		return options, err
	}
	headerEnd, err := regexp.Compile("^-->$")
	if err != nil {
		return options, err
	}
	version, err := regexp.Compile("\\s+goNetViz \"([0-9.]+)\"$")
	if err != nil {
		return options, err
	}

	for svg.Scan() {
		line := svg.Text()
		switch {
		case options.LimitX == 0 && options.LimitY == 0 && header == false:
			matches := limits.FindStringSubmatch(line)
			if len(matches) == 3 {
				options.LimitX, _ = strconv.Atoi(matches[1])
				options.LimitY, _ = strconv.Atoi(matches[2])
			}
		case header == false:
			if headerStart.MatchString(line) {
				header = true
			}
		case len(variant) == 0:
			matches := version.FindStringSubmatch(line)
			if len(matches) == 2 {
				variant, err = checkVersion(&parseOptions, matches[1])
				if err != nil {
					return options, err
				}
			}
		default:
			if optionIndex > len(parseOptions) {
				return options, fmt.Errorf("Option index is out of range")
			}
			regex, err := regexp.Compile(parseOptions[optionIndex].regex)
			if err != nil {
				return options, err
			}
			matches := regex.FindStringSubmatch(line)
			if len(matches) == 2 {
				option := reflect.ValueOf(&options).Elem().FieldByName(parseOptions[optionIndex].reconstructOption)

				switch option.Kind() {
				case reflect.Int:
					new, _ := strconv.Atoi(matches[1])
					option.SetInt(int64(new))
				case reflect.String:
					option.SetString(matches[1])
				default:
					return options, fmt.Errorf("Unhandeld option type")
				}
				optionIndex += 1
			} else {
				if headerEnd.MatchString(line) {
					return options, nil
				}
			}
		}
	}
	return options, fmt.Errorf("No end of header found")
}

func extractInformation(g *errgroup.Group, ch chan []byte, cfg configs) error {
	inputfile, err := os.Open(cfg.file)
	if err != nil {
		return fmt.Errorf("Could not open file %s: %s\n", cfg.file, err.Error())
	}
	defer inputfile.Close()
	svg := bufio.NewScanner(inputfile)
	var yLast int
	var packet []int
	defer close(ch)

	opt, err := checkHeader(svg)
	if err != nil {
		return err
	}

	pixel, err := regexp.Compile("^<rect x=\"(\\d+)\" y=\"(\\d+)\" width=\"\\d+\" height=\"\\d+\" style=\"fill:rgb\\((\\d+),(\\d+),(\\d+)\\)\" />$")
	if err != nil {
		return err
	}
	svgEnd, err := regexp.Compile("</svg>")
	if err != nil {
		return err
	}

	for svg.Scan() {
		line := svg.Text()
		matches := pixel.FindStringSubmatch(line)
		if len(matches) == 6 {
			pixelX, _ := strconv.Atoi(matches[1])
			pixelY, _ := strconv.Atoi(matches[2])
			if pixelY != yLast {
				yLast = pixelY
				if err := createPacket(ch, packet, opt.BpP); err != nil {
					return err
				}
				packet = packet[:0]
			}
			if pixelX >= opt.LimitX {
				return fmt.Errorf("x-coordinate (%d) is bigger than the limit (%d)\n", pixelX, opt.LimitX)
			}
			r, _ := strconv.Atoi(matches[3])
			g, _ := strconv.Atoi(matches[4])
			b, _ := strconv.Atoi(matches[5])
			packet = append(packet, r, g, b)
		} else if svgEnd.MatchString(line) {
			if len(packet) != 0 {
				return createPacket(ch, packet, opt.BpP)
			}
		}
	}
	return nil
}

func createPcap(g *errgroup.Group, ch chan []byte, cfg configs) error {
	filename := cfg.prefix
	filename += ".pcap"
	output, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("Could not create file %s: %s\n", filename, err.Error())
	}
	defer output.Close()
	w := pcapgo.NewWriter(output)
	w.WriteFileHeader(65536, layers.LinkTypeEthernet)

	for i, ok := <-ch; ok; i, ok = <-ch {
		w.WritePacket(gopacket.CaptureInfo{CaptureLength: len(i), Length: len(i), InterfaceIndex: 0}, i)
	}

	return nil
}

func reconstruct(g *errgroup.Group, cfg configs) error {
	ch := make(chan []byte)

	go extractInformation(g, ch, cfg)

	g.Go(func() error {
		return createPcap(g, ch, cfg)
	})

	return g.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
/*
 *  Filename:    dispatch.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description: Defines the Dispatch type and methods.
 */

//  Package dispatch provides goroutine dispatch and concurrency limiting.
//  It provides an object Dispatch which is a queueing system for concurrent
//  functions. It implements a dynamic limit on the number of routines that
//  it runs simultaneously. It also uses a Queue interface, allowing for
//  alternate queue implementations.
//
//  See github.com/bmatsuo/dispatch/queues for more about the Queue
//  interface and, a list of commonly used queues.
//
//  See github.com/bmatsuo/dispatch/examples for usage examples.
package dispatch

import (
    "sync"
    //"log"
    "github.com/bmatsuo/dispatch/queues"
)

//  A Dispatch is an automated function dispatch queue with a limited
//  number of concurrent gorountines. The queue can be altered with the
//  Dispatch methods Enqueue and SetKey (TODO add method Remove).
type Dispatch struct {
    // The maximum number of goroutines can be changed while the queue is
    // processing.
    MaxGo      int

    // Handle waiting when the limit of concurrent goroutines has been reached.
    waitingToRun bool
    nextWait     *sync.WaitGroup

    // Handle waiting when function queue is empty.
    waitingOnQ   bool
    restart      *sync.WaitGroup

    // Manage the Start()'ing of a Dispatch, avoiding race conditions.
    startLock    *sync.Mutex
    started      bool

    // Handle goroutine-safe queue operations.
    qLock        *sync.Mutex
    queue        queues.Queue

    // Handle goroutine-safe limiting and identifier operations.
    pLock        *sync.Mutex
    processing   int         // Number of QueueTasks running
    idcount      int64       // pid counter

    // The longest the dispatch queue grew.
    maxlength    int

    // Handle stopping of the Start() method.
    kill         chan bool
}

//  Create a new Dispatch object with a specified limit on concurrency.
func New(maxroutines int) *Dispatch {
    return NewCustom(maxroutines, queues.NewFIFO())
}

//  Create a new Dispatch object with a custom backend queue satisfying
//  interface queues.Queue. It is not safe to allow non-Dispatch methods
//  any access to the object queue. This can lead to race conditions with
//  possible corruption of internal structures. So, it's considered a best
//  practice to only pass NewCustom(...) newly created queues.
//      fifoDispatch     := NewCustom(10, queues.NewFIFO())
//      priorityDispatch := NewCustom(20, queues.NewPriorityQueue())
func NewCustom(maxroutines int, queue queues.Queue) *Dispatch {
    var d = new(Dispatch)
    d.startLock = new(sync.Mutex)
    d.qLock     = new(sync.Mutex)
    d.pLock     = new(sync.Mutex)
    d.restart   = new(sync.WaitGroup)
    d.kill      = make(chan bool)
    d.nextWait  = new(sync.WaitGroup)
    d.queue     = queue
    d.MaxGo     = maxroutines
    d.idcount   = 0
    d.maxlength = 0
    return d
}

//  A simple task for use in a priority-less queue. See package
//  github.com/bmatsuo/dispatch/queues
type StdTask struct {
    F func(id int64)
}

func NewTask(f func(int64)) *StdTask {
    t := new(StdTask)
    t.F = f
    return t
}
func (dt *StdTask) Type() string {
    return "StdTask"
}
func (dt *StdTask) SetFunc(f func(id int64)) {
    dt.F = f
}
func (dt *StdTask) Func() func(id int64) {
    return dt.F
}

type dispatchTaskWrapper struct {
    id  int64
    t   queues.Task
}

func (dtw dispatchTaskWrapper) Func() func(id int64) {
    return dtw.t.Func()
}
func (dtw dispatchTaskWrapper) Id() int64 {
    return dtw.id
}
func (dtw dispatchTaskWrapper) Task() queues.Task {
    return dtw.t
}

func (gq *Dispatch) Len() int {
    return gq.queue.Len()
}
func (gq *Dispatch) MaxLen() int {
    return gq.maxlength
}
//  Enqueue a task for execution as a goroutine. The given queues.Task is
//  given a unique id (int64) and stored in the Dispatch gq's backend
//  queues.Queue object.
func (gq *Dispatch) Enqueue(t queues.Task) int64 {
    // Wrap the function so it works with the goroutine limiting code.
    var f = t.Func()
    var dtFunc = func (id int64) {
        // Run the given function.
        f(id)

        // Decrement the process counter.
        gq.pLock.Lock()
        //log.Printf("processing: %d, waiting: %v", gq.processing, gq.waitingToRun)
        gq.processing--
        if gq.waitingToRun {
            gq.waitingToRun = false
            gq.nextWait.Done()
        }
        gq.pLock.Unlock()
    }
    t.SetFunc(dtFunc)

    // Lock the queue and enqueue a new task.
    gq.qLock.Lock()
    gq.idcount++
    var id = gq.idcount
    gq.queue.Enqueue(dispatchTaskWrapper{id, t})
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.queue.Len() > gq.maxlength {
        gq.maxlength = gq.queue.Len()
    }
    gq.qLock.Unlock()

    return id
}

//  Stop the queue after gq.Start() has been called. Any goroutines which
//  have not already been dequeued will not be executed until gq.Start()
//  is called again.
func (gq *Dispatch) Stop() {
    // Lock out Start() and queue ops for the entire call.
    gq.startLock.Lock()
    defer gq.startLock.Unlock()
    gq.qLock.Lock()
    defer gq.qLock.Unlock()

    if !gq.started {
        return
    }

    // Clear channel flags and close channels, stoping further processing.
    close(gq.kill)
    gq.started = false
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.waitingToRun {
        gq.waitingToRun = false
        gq.nextWait.Done()
    }
}

//  Start the next task in the queue. It's assumed that the queue is non-
//  empty. Furthermore, there should only be one goroutine in this method
//  (for this object) at a time. Both conditions are enforced in
//  gq.Start(), which calls gq.next() exclusively.
func (gq *Dispatch) next() {
    for true {
        // Attempt to start processing the file.
        gq.pLock.Lock()
        if gq.processing >= gq.MaxGo {
            gq.waitingToRun = true
            gq.nextWait.Add(1)
            gq.pLock.Unlock()
            gq.nextWait.Wait()
            continue
        }
        // Keep the books and reset wait time before unlocking.
        gq.processing++
        gq.pLock.Unlock()

        // Get an element from the queue.
        gq.qLock.Lock()
        var wrapper = gq.queue.Dequeue().(queues.RegisteredTask)
        gq.qLock.Unlock()

        // Begin processing and asyncronously return.
        //var task = taskelm.Value.(dispatchTaskWrapper)
        var task = wrapper.Func()
        go task(wrapper.Id())
        return
    }
}

//  Start executing goroutines. Don't stop until gq.Stop() is called. This
//  method will take control of the calling thread. But, it's safe to call
//  in a goroutine.
//      gq := dispatch.New()
//      go gq.Start()
//      wg := new(sync.WaitGroup)
//      for i := 0 ; i < 1000 ; i++ {
//          wg.Add(1)
//          gq.Enqueue(&dispatch.StdTask{
//              func(id int64) {
//                  log.Printf("I'm alive %d", i)
//                  wg.Done()
//              } } )
//      }
//      wg.Wait()
//      gq.Stop()
func (gq *Dispatch) Start() {
    // Avoid multiple gq.Start() methods and avoid race conditions.
    gq.startLock.Lock()
    if gq.started {
        panic("already started")
    }
    gq.started = true
    gq.startLock.Unlock()

    // Recreate any channels that were closed by a previous Stop().
    var inited = false
    for !inited {
        select {
        case _, okKill := <-gq.kill:
            if !okKill {
                gq.kill = make(chan bool)
            }
        default:
            inited = true
        }
    }

    // Process the queue
    for true {
        select {
        case die, ok := <-gq.kill:
            // If something came out of this channel, we must stop.
            if !ok {
                // Recreate the channel.
                gq.kill = make(chan bool)
                return
            }
            if die {
                return
            }
        default:
            // Check the queue size and determine if we need to wait.
            gq.qLock.Lock()
            var wait = gq.queue.Len() == 0
            if gq.waitingOnQ = wait; wait {
                gq.restart.Add(1)
            }
            gq.qLock.Unlock()

            if wait {
                // Wait for a restart signal from gq.Enqueue
                gq.restart.Wait()
            } else {
                // Process the head of the queue and start the loop again.
                gq.next()
                continue
            }
        }
    }
}
<commit_msg>Update documentation examples with the StdTask ctor.<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
/*
 *  Filename:    dispatch.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description: Defines the Dispatch type and methods.
 */

//  Package dispatch provides goroutine dispatch and concurrency limiting.
//  It provides an object Dispatch which is a queueing system for concurrent
//  functions. It implements a dynamic limit on the number of routines that
//  it runs simultaneously. It also uses a Queue interface, allowing for
//  alternate queue implementations.
//
//  See github.com/bmatsuo/dispatch/queues for more about the Queue
//  interface and, a list of commonly used queues.
//
//  See github.com/bmatsuo/dispatch/examples for usage examples.
package dispatch

import (
    "sync"
    //"log"
    "github.com/bmatsuo/dispatch/queues"
)

//  A Dispatch is an automated function dispatch queue with a limited
//  number of concurrent gorountines. The queue can be altered with the
//  Dispatch methods Enqueue and SetKey (TODO add method Remove).
type Dispatch struct {
    // The maximum number of goroutines can be changed while the queue is
    // processing.
    MaxGo      int

    // Handle waiting when the limit of concurrent goroutines has been reached.
    waitingToRun bool
    nextWait     *sync.WaitGroup

    // Handle waiting when function queue is empty.
    waitingOnQ   bool
    restart      *sync.WaitGroup

    // Manage the Start()'ing of a Dispatch, avoiding race conditions.
    startLock    *sync.Mutex
    started      bool

    // Handle goroutine-safe queue operations.
    qLock        *sync.Mutex
    queue        queues.Queue

    // Handle goroutine-safe limiting and identifier operations.
    pLock        *sync.Mutex
    processing   int         // Number of QueueTasks running
    idcount      int64       // pid counter

    // The longest the dispatch queue grew.
    maxlength    int

    // Handle stopping of the Start() method.
    kill         chan bool
}

//  Create a new Dispatch object with a specified limit on concurrency.
func New(maxroutines int) *Dispatch {
    return NewCustom(maxroutines, queues.NewFIFO())
}

//  Create a new Dispatch object with a custom backend queue satisfying
//  interface queues.Queue. It is not safe to allow non-Dispatch methods
//  any access to the object queue. This can lead to race conditions with
//  possible corruption of internal structures. So, it's considered a best
//  practice to only pass NewCustom(...) newly created queues.
//      fifoDispatch     := NewCustom(10, queues.NewFIFO())
//      priorityDispatch := NewCustom(20, queues.NewPriorityQueue())
func NewCustom(maxroutines int, queue queues.Queue) *Dispatch {
    var d = new(Dispatch)
    d.startLock = new(sync.Mutex)
    d.qLock     = new(sync.Mutex)
    d.pLock     = new(sync.Mutex)
    d.restart   = new(sync.WaitGroup)
    d.kill      = make(chan bool)
    d.nextWait  = new(sync.WaitGroup)
    d.queue     = queue
    d.MaxGo     = maxroutines
    d.idcount   = 0
    d.maxlength = 0
    return d
}

//  A simple task for use in a priority-less queue. See package
//  github.com/bmatsuo/dispatch/queues
type StdTask struct {
    F func(id int64)
}

func NewTask(f func(int64)) *StdTask {
    t := new(StdTask)
    t.F = f
    return t
}
func (dt *StdTask) Type() string {
    return "StdTask"
}
func (dt *StdTask) SetFunc(f func(id int64)) {
    dt.F = f
}
func (dt *StdTask) Func() func(id int64) {
    return dt.F
}

type dispatchTaskWrapper struct {
    id  int64
    t   queues.Task
}

func (dtw dispatchTaskWrapper) Func() func(id int64) {
    return dtw.t.Func()
}
func (dtw dispatchTaskWrapper) Id() int64 {
    return dtw.id
}
func (dtw dispatchTaskWrapper) Task() queues.Task {
    return dtw.t
}

func (gq *Dispatch) Len() int {
    return gq.queue.Len()
}
func (gq *Dispatch) MaxLen() int {
    return gq.maxlength
}
//  Enqueue a task for execution as a goroutine. The given queues.Task is
//  given a unique id (int64) and stored in the Dispatch gq's backend
//  queues.Queue object.
func (gq *Dispatch) Enqueue(t queues.Task) int64 {
    // Wrap the function so it works with the goroutine limiting code.
    var f = t.Func()
    var dtFunc = func (id int64) {
        // Run the given function.
        f(id)

        // Decrement the process counter.
        gq.pLock.Lock()
        //log.Printf("processing: %d, waiting: %v", gq.processing, gq.waitingToRun)
        gq.processing--
        if gq.waitingToRun {
            gq.waitingToRun = false
            gq.nextWait.Done()
        }
        gq.pLock.Unlock()
    }
    t.SetFunc(dtFunc)

    // Lock the queue and enqueue a new task.
    gq.qLock.Lock()
    gq.idcount++
    var id = gq.idcount
    gq.queue.Enqueue(dispatchTaskWrapper{id, t})
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.queue.Len() > gq.maxlength {
        gq.maxlength = gq.queue.Len()
    }
    gq.qLock.Unlock()

    return id
}

//  Stop the queue after gq.Start() has been called. Any goroutines which
//  have not already been dequeued will not be executed until gq.Start()
//  is called again.
func (gq *Dispatch) Stop() {
    // Lock out Start() and queue ops for the entire call.
    gq.startLock.Lock()
    defer gq.startLock.Unlock()
    gq.qLock.Lock()
    defer gq.qLock.Unlock()

    if !gq.started {
        return
    }

    // Clear channel flags and close channels, stoping further processing.
    close(gq.kill)
    gq.started = false
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.waitingToRun {
        gq.waitingToRun = false
        gq.nextWait.Done()
    }
}

//  Start the next task in the queue. It's assumed that the queue is non-
//  empty. Furthermore, there should only be one goroutine in this method
//  (for this object) at a time. Both conditions are enforced in
//  gq.Start(), which calls gq.next() exclusively.
func (gq *Dispatch) next() {
    for true {
        // Attempt to start processing the file.
        gq.pLock.Lock()
        if gq.processing >= gq.MaxGo {
            gq.waitingToRun = true
            gq.nextWait.Add(1)
            gq.pLock.Unlock()
            gq.nextWait.Wait()
            continue
        }
        // Keep the books and reset wait time before unlocking.
        gq.processing++
        gq.pLock.Unlock()

        // Get an element from the queue.
        gq.qLock.Lock()
        var wrapper = gq.queue.Dequeue().(queues.RegisteredTask)
        gq.qLock.Unlock()

        // Begin processing and asyncronously return.
        //var task = taskelm.Value.(dispatchTaskWrapper)
        var task = wrapper.Func()
        go task(wrapper.Id())
        return
    }
}

//  Start executing goroutines. Don't stop until gq.Stop() is called. This
//  method will take control of the calling thread. But, it's safe to call
//  in a goroutine.
//      gq := dispatch.New()
//      go gq.Start()
//      wg := new(sync.WaitGroup)
//      for i := 0 ; i < 1000 ; i++ {
//          wg.Add(1)
//          gq.Enqueue(dispatch.NewTask(
//              func(id int64) {
//                  log.Printf("I'm alive %d", i)
//                  wg.Done()
//              } ) )
//      }
//      wg.Wait()
//      gq.Stop()
func (gq *Dispatch) Start() {
    // Avoid multiple gq.Start() methods and avoid race conditions.
    gq.startLock.Lock()
    if gq.started {
        panic("already started")
    }
    gq.started = true
    gq.startLock.Unlock()

    // Recreate any channels that were closed by a previous Stop().
    var inited = false
    for !inited {
        select {
        case _, okKill := <-gq.kill:
            if !okKill {
                gq.kill = make(chan bool)
            }
        default:
            inited = true
        }
    }

    // Process the queue
    for true {
        select {
        case die, ok := <-gq.kill:
            // If something came out of this channel, we must stop.
            if !ok {
                // Recreate the channel.
                gq.kill = make(chan bool)
                return
            }
            if die {
                return
            }
        default:
            // Check the queue size and determine if we need to wait.
            gq.qLock.Lock()
            var wait = gq.queue.Len() == 0
            if gq.waitingOnQ = wait; wait {
                gq.restart.Add(1)
            }
            gq.qLock.Unlock()

            if wait {
                // Wait for a restart signal from gq.Enqueue
                gq.restart.Wait()
            } else {
                // Process the head of the queue and start the loop again.
                gq.next()
                continue
            }
        }
    }
}
<|endoftext|>


<commit_before>package app

import (
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/fragmenta/server"
	"github.com/fragmenta/server/config"
	"github.com/fragmenta/server/log"
	"github.com/fragmenta/view"
)

// Serve static files (assets, images etc)
func fileHandler(w http.ResponseWriter, r *http.Request) error {

	// First try serving assets
	err := serveAsset(w, r)
	if err == nil {
		return nil
	}

	// If assets fail, try to serve file in public
	return serveFile(w, r)
}

// serveFile serves a file from ./public if it exists
func serveFile(w http.ResponseWriter, r *http.Request) error {

	// Try a local path in the public directory
	localPath := "./public" + path.Clean(r.URL.Path)
	s, err := os.Stat(localPath)
	if err != nil {
		// If file not found return 404
		if os.IsNotExist(err) {
			return server.NotFoundError(err)
		}

		// For other errors return not authorised
		return server.NotAuthorizedError(err)
	}

	// If not a file return immediately
	if s.IsDir() {
		return nil
	}

	// If the file exists and we can access it, serve it with cache control
	w.Header().Set("Cache-Control", "max-age:3456000, public")
	http.ServeFile(w, r, localPath)
	return nil
}

// serveAsset serves a file from ./public/assets usings appAssets
func serveAsset(w http.ResponseWriter, r *http.Request) error {

	p := path.Clean(r.URL.Path)

	// It must be under /assets, or we don't serve
	if !strings.HasPrefix(p, "/assets/") {
		return server.NotFoundError(nil)
	}

	// Try to find an asset in our list
	f := appAssets.File(path.Base(p))
	if f == nil {
		return server.NotFoundError(nil)
	}

	// Serve the local file, with cache control
	localPath := "./" + f.LocalPath()
	w.Header().Set("Cache-Control", "max-age:3456000, public")
	http.ServeFile(w, r, localPath)
	return nil
}

// errHandler renders an error using error templates if available
func errHandler(w http.ResponseWriter, r *http.Request, e error) {

	// Cast the error to a status error if it is one, if not wrap it in a Status 500 error
	err := server.ToStatusError(e)
	log.Error(log.V{"error": err})

	view := view.NewWithPath("", w)
	view.AddKey("title", err.Title)
	view.AddKey("message", err.Message)
	// In production, provide no detail for security reasons
	if !config.Production() {
		view.AddKey("status", err.Status)
		view.AddKey("file", err.FileLine())
		view.AddKey("error", err.Err)
	}
	view.Template("app/views/error.html.got")
	w.WriteHeader(err.Status)
	view.Render()
}
<commit_msg>Add etag for assets<commit_after>package app

import (
	"crypto/sha1"
	"encoding/hex"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/fragmenta/server"
	"github.com/fragmenta/server/config"
	"github.com/fragmenta/server/log"
	"github.com/fragmenta/view"
)

// Serve static files (assets, images etc)
func fileHandler(w http.ResponseWriter, r *http.Request) error {

	// First try serving assets
	err := serveAsset(w, r)
	if err == nil {
		return nil
	}

	// If assets fail, try to serve file in public
	return serveFile(w, r)
}

// serveFile serves a file from ./public if it exists
func serveFile(w http.ResponseWriter, r *http.Request) error {

	// Try a local path in the public directory
	localPath := "./public" + path.Clean(r.URL.Path)
	s, err := os.Stat(localPath)
	if err != nil {
		// If file not found return 404
		if os.IsNotExist(err) {
			return server.NotFoundError(err)
		}

		// For other errors return not authorised
		return server.NotAuthorizedError(err)
	}

	// If not a file return immediately
	if s.IsDir() {
		return nil
	}

	// If the file exists and we can access it, serve it with cache control
	// if in production
	if config.Production() {
		// Cache for 30 days
		w.Header().Set("Cache-Control", "max-age:2592000, public")
		// For etag Just hash the path - static resources are assumed to have a fingerprint
		w.Header().Set("ETag", hash(r.URL.Path))
	}

	http.ServeFile(w, r, localPath)
	return nil
}

// hash returns the sha hash of a string
func hash(s string) string {
	sum := sha1.Sum([]byte(s))
	return hex.EncodeToString([]byte(sum[:]))
}

// serveAsset serves a file from ./public/assets usings appAssets
func serveAsset(w http.ResponseWriter, r *http.Request) error {

	p := path.Clean(r.URL.Path)

	// It must be under /assets, or we don't serve
	if !strings.HasPrefix(p, "/assets/") {
		return server.NotFoundError(nil)
	}

	// Try to find an asset in our list
	f := appAssets.File(path.Base(p))
	if f == nil {
		return server.NotFoundError(nil)
	}

	// Serve the local file, with cache control
	localPath := "./" + f.LocalPath()
	// If the file exists and we can access it, serve it with cache control in production
	if config.Production() {
		// Cache for 30 days
		w.Header().Set("Cache-Control", "max-age:2592000, public")
		// For etag Just hash the path - static resources are assumed to have a fingerprint
		w.Header().Set("ETag", hash(r.URL.Path))
	}
	http.ServeFile(w, r, localPath)
	return nil
}

// errHandler renders an error using error templates if available
func errHandler(w http.ResponseWriter, r *http.Request, e error) {

	// Cast the error to a status error if it is one, if not wrap it in a Status 500 error
	err := server.ToStatusError(e)
	log.Error(log.V{"error": err})

	view := view.NewWithPath("", w)
	view.AddKey("title", err.Title)
	view.AddKey("message", err.Message)
	// In production, provide no detail for security reasons
	if !config.Production() {
		view.AddKey("status", err.Status)
		view.AddKey("file", err.FileLine())
		view.AddKey("error", err.Err)
	}
	view.Template("app/views/error.html.got")
	w.WriteHeader(err.Status)
	view.Render()
}
<|endoftext|>


<commit_before>package check

import (
	"fmt"
	"plaid/parser"
	"plaid/types"
	"plaid/vm"
)

var binops = map[string]map[types.Type]map[types.Type]types.Type{
	"+": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
		types.Str: map[types.Type]types.Type{types.Str: types.Str},
	},
	"-": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"*": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"<": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"<=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"[": map[types.Type]map[types.Type]types.Type{
		types.Str: map[types.Type]types.Type{types.Int: types.Str},
	},
}

// Check takes an existing abstract syntax tree and performs type checks and
// other correctness checks. It returns a list of any errors that were
// discovered inside the AST
func Check(prog parser.Program, libraries ...vm.Library) *Scope {
	global := makeScope(nil, nil)

	for _, library := range libraries {
		for name, builtin := range library {
			global.registerLocalVariable(name, builtin.Type)
		}
	}

	checkProgram(global, prog)
	return global
}

func checkProgram(scope *Scope, prog parser.Program) {
	for _, stmt := range prog.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkStmt(scope *Scope, stmt parser.Stmt) {
	switch stmt := stmt.(type) {
	case parser.DeclarationStmt:
		checkDeclarationStmt(scope, stmt)
		break
	case parser.ReturnStmt:
		checkReturnStmt(scope, stmt)
		break
	case parser.ExprStmt:
		checkExprAllowVoid(scope, stmt.Expr)
		break
	}

	if scope.hasBodyQueue() {
		sig, expr := scope.dequeueBody()
		checkFunctionBody(scope, sig, expr)
	}
}

func checkStmtBlock(scope *Scope, block parser.StmtBlock) {
	for _, stmt := range block.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkDeclarationStmt(scope *Scope, stmt parser.DeclarationStmt) {
	name := stmt.Name.Name
	typ := checkExpr(scope, stmt.Expr)
	scope.registerLocalVariable(name, typ)
}

func checkReturnStmt(scope *Scope, stmt parser.ReturnStmt) {
	var ret types.Type = types.TypeVoid{}
	if stmt.Expr != nil {
		ret = checkExpr(scope, stmt.Expr)
	}

	if scope.pendingReturn == nil {
		scope.addError(fmt.Errorf("return statements must be inside a function"))
		return
	}

	if scope.pendingReturn.Equals(ret) || ret.IsError() {
		return
	}

	if scope.pendingReturn.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected to return nothing, got '%s'", ret))
		return
	}

	if ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected a return type of '%s', got nothing", scope.pendingReturn))
	}

	scope.addError(fmt.Errorf("expected to return '%s', got '%s'", scope.pendingReturn, ret))
}

func checkExprAllowVoid(scope *Scope, expr parser.Expr) types.Type {
	var typ types.Type = types.TypeError{}
	switch expr := expr.(type) {
	case parser.FunctionExpr:
		typ = checkFunctionExpr(scope, expr)
	case parser.DispatchExpr:
		typ = checkDispatchExpr(scope, expr)
	case parser.AssignExpr:
		typ = checkAssignExpr(scope, expr)
	case parser.BinaryExpr:
		typ = checkBinaryExpr(scope, expr)
	case parser.SubscriptExpr:
		typ = checkSubscriptExpr(scope, expr)
	case parser.IdentExpr:
		typ = checkIdentExpr(scope, expr)
	case parser.NumberExpr:
		typ = checkNumberExpr(scope, expr)
	case parser.StringExpr:
		typ = checkStringExpr(scope, expr)
	case parser.BooleanExpr:
		typ = checkBooleanExpr(scope, expr)
	default:
		scope.addError(fmt.Errorf("unknown expression type"))
	}

	return typ
}

func checkExpr(scope *Scope, expr parser.Expr) types.Type {
	typ := checkExprAllowVoid(scope, expr)

	if typ.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("cannot use void types in an expression"))
		return types.TypeError{}
	}

	return typ
}

func checkFunctionExpr(scope *Scope, expr parser.FunctionExpr) types.Type {
	ret := types.ConvertTypeNote(expr.Ret)
	params := []types.Type{}
	for _, param := range expr.Params {
		params = append(params, types.ConvertTypeNote(param.Note))
	}
	tuple := types.TypeTuple{Children: params}

	sig := types.TypeFunction{Params: tuple, Ret: ret}
	scope.enqueueBody(ret, expr)
	return sig
}

func checkFunctionBody(scope *Scope, ret types.Type, expr parser.FunctionExpr) {
	pushed := makeScope(scope, ret)

	for _, param := range expr.Params {
		paramName := param.Name.Name
		paramType := types.ConvertTypeNote(param.Note)
		pushed.registerLocalVariable(paramName, paramType)
	}

	checkStmtBlock(pushed, expr.Block)
}

func checkDispatchExpr(scope *Scope, expr parser.DispatchExpr) types.Type {
	// Resolve arguments to types
	argTypes := []types.Type{}
	for _, argExpr := range expr.Args {
		argTypes = append(argTypes, checkExpr(scope, argExpr))
	}

	// Resolve callee to type
	calleeType := checkExpr(scope, expr.Callee)
	calleeFunc, ok := calleeType.(types.TypeFunction)
	if ok == false {
		if calleeType.IsError() == false {
			scope.addError(fmt.Errorf("cannot call function on type '%s'", calleeType))
		}

		return types.TypeError{}
	}

	// Resolve return type
	retType := calleeFunc.Ret

	// Check that the given argument types match the expected parameter types
	totalArgs := len(argTypes)
	totalParams := len(calleeFunc.Params.Children)
	if totalArgs == totalParams {
		for i := 0; i < totalArgs; i++ {
			argType := argTypes[i]
			paramType := calleeFunc.Params.Children[i]

			if argType.Equals(paramType) == false {
				scope.addError(fmt.Errorf("expected '%s', got '%s'", paramType, argType))
				retType = types.TypeError{}
			}
		}
	} else {
		scope.addError(fmt.Errorf("expected %d arguments, got %d", totalParams, totalArgs))
		retType = types.TypeError{}
	}

	return retType
}

func checkAssignExpr(scope *Scope, expr parser.AssignExpr) types.Type {
	name := expr.Left.Name
	leftType := scope.getVariable(name)
	rightType := checkExpr(scope, expr.Right)

	if leftType == nil {
		scope.addError(fmt.Errorf("'%s' cannot be assigned before it is declared", name))
		return types.TypeError{}
	}

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if leftType.Equals(rightType) == false {
		scope.addError(fmt.Errorf("'%s' cannot be assigned type '%s'", leftType, rightType))
		return types.TypeError{}
	}

	return leftType
}

func checkBinaryExpr(scope *Scope, expr parser.BinaryExpr) types.Type {
	leftType := checkExpr(scope, expr.Left)
	rightType := checkExpr(scope, expr.Right)

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if operLUT, ok := binops[expr.Oper]; ok {
		if leftLUT, ok := operLUT[leftType]; ok {
			if retType, ok := leftLUT[rightType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("operator '%s' does not support %s and %s", expr.Oper, leftType, rightType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '%s'", expr.Oper))
	return types.TypeError{}
}

func checkSubscriptExpr(scope *Scope, expr parser.SubscriptExpr) types.Type {
	listType := checkExpr(scope, expr.ListLike)
	indexType := checkExpr(scope, expr.Index)

	if listType.IsError() || indexType.IsError() {
		return types.TypeError{}
	}

	if subscriptLUT, ok := binops["["]; ok {
		if listLUT, ok := subscriptLUT[listType]; ok {
			if retType, ok := listLUT[indexType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("subscript operator does not support %s[%s]", listType, indexType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '['"))
	return types.TypeError{}
}

func checkIdentExpr(scope *Scope, expr parser.IdentExpr) types.Type {
	if scope.existingVariable(expr.Name) {
		return scope.getVariable(expr.Name)
	}

	scope.addError(fmt.Errorf("variable '%s' was used before it was declared", expr.Name))
	return types.TypeError{}
}

func checkNumberExpr(scope *Scope, expr parser.NumberExpr) types.Type {
	return types.Int
}

func checkStringExpr(scope *Scope, expr parser.StringExpr) types.Type {
	return types.Str
}

func checkBooleanExpr(scope *Scope, expr parser.BooleanExpr) types.Type {
	return types.Bool
}
<commit_msg>typecheck if statement<commit_after>package check

import (
	"fmt"
	"plaid/parser"
	"plaid/types"
	"plaid/vm"
)

var binops = map[string]map[types.Type]map[types.Type]types.Type{
	"+": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
		types.Str: map[types.Type]types.Type{types.Str: types.Str},
	},
	"-": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"*": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"<": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"<=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"[": map[types.Type]map[types.Type]types.Type{
		types.Str: map[types.Type]types.Type{types.Int: types.Str},
	},
}

// Check takes an existing abstract syntax tree and performs type checks and
// other correctness checks. It returns a list of any errors that were
// discovered inside the AST
func Check(prog parser.Program, libraries ...vm.Library) *Scope {
	global := makeScope(nil, nil)

	for _, library := range libraries {
		for name, builtin := range library {
			global.registerLocalVariable(name, builtin.Type)
		}
	}

	checkProgram(global, prog)
	return global
}

func checkProgram(scope *Scope, prog parser.Program) {
	for _, stmt := range prog.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkStmt(scope *Scope, stmt parser.Stmt) {
	switch stmt := stmt.(type) {
	case parser.IfStmt:
		checkIfStmt(scope, stmt)
		break
	case parser.DeclarationStmt:
		checkDeclarationStmt(scope, stmt)
		break
	case parser.ReturnStmt:
		checkReturnStmt(scope, stmt)
		break
	case parser.ExprStmt:
		checkExprAllowVoid(scope, stmt.Expr)
		break
	}

	if scope.hasBodyQueue() {
		sig, expr := scope.dequeueBody()
		checkFunctionBody(scope, sig, expr)
	}
}

func checkStmtBlock(scope *Scope, block parser.StmtBlock) {
	for _, stmt := range block.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkIfStmt(scope *Scope, stmt parser.IfStmt) {
	typ := checkExpr(scope, stmt.Cond)
	if typ.Equals(types.Bool) == false {
		scope.addError(fmt.Errorf("condition must resolve to a boolean"))
	}

	checkStmtBlock(scope, stmt.Clause)
}

func checkDeclarationStmt(scope *Scope, stmt parser.DeclarationStmt) {
	name := stmt.Name.Name
	typ := checkExpr(scope, stmt.Expr)
	scope.registerLocalVariable(name, typ)
}

func checkReturnStmt(scope *Scope, stmt parser.ReturnStmt) {
	var ret types.Type = types.TypeVoid{}
	if stmt.Expr != nil {
		ret = checkExpr(scope, stmt.Expr)
	}

	if scope.pendingReturn == nil {
		scope.addError(fmt.Errorf("return statements must be inside a function"))
		return
	}

	if scope.pendingReturn.Equals(ret) || ret.IsError() {
		return
	}

	if scope.pendingReturn.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected to return nothing, got '%s'", ret))
		return
	}

	if ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected a return type of '%s', got nothing", scope.pendingReturn))
	}

	scope.addError(fmt.Errorf("expected to return '%s', got '%s'", scope.pendingReturn, ret))
}

func checkExprAllowVoid(scope *Scope, expr parser.Expr) types.Type {
	var typ types.Type = types.TypeError{}
	switch expr := expr.(type) {
	case parser.FunctionExpr:
		typ = checkFunctionExpr(scope, expr)
	case parser.DispatchExpr:
		typ = checkDispatchExpr(scope, expr)
	case parser.AssignExpr:
		typ = checkAssignExpr(scope, expr)
	case parser.BinaryExpr:
		typ = checkBinaryExpr(scope, expr)
	case parser.SubscriptExpr:
		typ = checkSubscriptExpr(scope, expr)
	case parser.IdentExpr:
		typ = checkIdentExpr(scope, expr)
	case parser.NumberExpr:
		typ = checkNumberExpr(scope, expr)
	case parser.StringExpr:
		typ = checkStringExpr(scope, expr)
	case parser.BooleanExpr:
		typ = checkBooleanExpr(scope, expr)
	default:
		scope.addError(fmt.Errorf("unknown expression type"))
	}

	return typ
}

func checkExpr(scope *Scope, expr parser.Expr) types.Type {
	typ := checkExprAllowVoid(scope, expr)

	if typ.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("cannot use void types in an expression"))
		return types.TypeError{}
	}

	return typ
}

func checkFunctionExpr(scope *Scope, expr parser.FunctionExpr) types.Type {
	ret := types.ConvertTypeNote(expr.Ret)
	params := []types.Type{}
	for _, param := range expr.Params {
		params = append(params, types.ConvertTypeNote(param.Note))
	}
	tuple := types.TypeTuple{Children: params}

	sig := types.TypeFunction{Params: tuple, Ret: ret}
	scope.enqueueBody(ret, expr)
	return sig
}

func checkFunctionBody(scope *Scope, ret types.Type, expr parser.FunctionExpr) {
	pushed := makeScope(scope, ret)

	for _, param := range expr.Params {
		paramName := param.Name.Name
		paramType := types.ConvertTypeNote(param.Note)
		pushed.registerLocalVariable(paramName, paramType)
	}

	checkStmtBlock(pushed, expr.Block)
}

func checkDispatchExpr(scope *Scope, expr parser.DispatchExpr) types.Type {
	// Resolve arguments to types
	argTypes := []types.Type{}
	for _, argExpr := range expr.Args {
		argTypes = append(argTypes, checkExpr(scope, argExpr))
	}

	// Resolve callee to type
	calleeType := checkExpr(scope, expr.Callee)
	calleeFunc, ok := calleeType.(types.TypeFunction)
	if ok == false {
		if calleeType.IsError() == false {
			scope.addError(fmt.Errorf("cannot call function on type '%s'", calleeType))
		}

		return types.TypeError{}
	}

	// Resolve return type
	retType := calleeFunc.Ret

	// Check that the given argument types match the expected parameter types
	totalArgs := len(argTypes)
	totalParams := len(calleeFunc.Params.Children)
	if totalArgs == totalParams {
		for i := 0; i < totalArgs; i++ {
			argType := argTypes[i]
			paramType := calleeFunc.Params.Children[i]

			if argType.Equals(paramType) == false {
				scope.addError(fmt.Errorf("expected '%s', got '%s'", paramType, argType))
				retType = types.TypeError{}
			}
		}
	} else {
		scope.addError(fmt.Errorf("expected %d arguments, got %d", totalParams, totalArgs))
		retType = types.TypeError{}
	}

	return retType
}

func checkAssignExpr(scope *Scope, expr parser.AssignExpr) types.Type {
	name := expr.Left.Name
	leftType := scope.getVariable(name)
	rightType := checkExpr(scope, expr.Right)

	if leftType == nil {
		scope.addError(fmt.Errorf("'%s' cannot be assigned before it is declared", name))
		return types.TypeError{}
	}

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if leftType.Equals(rightType) == false {
		scope.addError(fmt.Errorf("'%s' cannot be assigned type '%s'", leftType, rightType))
		return types.TypeError{}
	}

	return leftType
}

func checkBinaryExpr(scope *Scope, expr parser.BinaryExpr) types.Type {
	leftType := checkExpr(scope, expr.Left)
	rightType := checkExpr(scope, expr.Right)

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if operLUT, ok := binops[expr.Oper]; ok {
		if leftLUT, ok := operLUT[leftType]; ok {
			if retType, ok := leftLUT[rightType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("operator '%s' does not support %s and %s", expr.Oper, leftType, rightType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '%s'", expr.Oper))
	return types.TypeError{}
}

func checkSubscriptExpr(scope *Scope, expr parser.SubscriptExpr) types.Type {
	listType := checkExpr(scope, expr.ListLike)
	indexType := checkExpr(scope, expr.Index)

	if listType.IsError() || indexType.IsError() {
		return types.TypeError{}
	}

	if subscriptLUT, ok := binops["["]; ok {
		if listLUT, ok := subscriptLUT[listType]; ok {
			if retType, ok := listLUT[indexType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("subscript operator does not support %s[%s]", listType, indexType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '['"))
	return types.TypeError{}
}

func checkIdentExpr(scope *Scope, expr parser.IdentExpr) types.Type {
	if scope.existingVariable(expr.Name) {
		return scope.getVariable(expr.Name)
	}

	scope.addError(fmt.Errorf("variable '%s' was used before it was declared", expr.Name))
	return types.TypeError{}
}

func checkNumberExpr(scope *Scope, expr parser.NumberExpr) types.Type {
	return types.Int
}

func checkStringExpr(scope *Scope, expr parser.StringExpr) types.Type {
	return types.Str
}

func checkBooleanExpr(scope *Scope, expr parser.BooleanExpr) types.Type {
	return types.Bool
}
<|endoftext|>


<commit_before>package check

import (
	"fmt"
	"plaid/parser"
	"plaid/types"
	"plaid/vm"
)

var binops = map[string]map[types.Type]map[types.Type]types.Type{
	"+": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
		types.Str: map[types.Type]types.Type{types.Str: types.Str},
	},
	"-": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"*": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Int},
	},
	"<": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"<=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	">=": map[types.Type]map[types.Type]types.Type{
		types.Int: map[types.Type]types.Type{types.Int: types.Bool},
	},
	"[": map[types.Type]map[types.Type]types.Type{
		types.Str: map[types.Type]types.Type{types.Int: types.Str},
	},
}

// Check takes an existing abstract syntax tree and performs type checks and
// other correctness checks. It returns a list of any errors that were
// discovered inside the AST
func Check(prog parser.Program, libraries ...vm.Library) *Scope {
	global := makeScope(nil, nil)

	for _, library := range libraries {
		for name, builtin := range library {
			global.registerLocalVariable(name, builtin.Type)
		}
	}

	checkProgram(global, prog)
	return global
}

func checkProgram(scope *Scope, prog parser.Program) {
	for _, stmt := range prog.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkStmt(scope *Scope, stmt parser.Stmt) {
	switch stmt := stmt.(type) {
	case parser.IfStmt:
		checkIfStmt(scope, stmt)
		break
	case parser.DeclarationStmt:
		checkDeclarationStmt(scope, stmt)
		break
	case parser.ReturnStmt:
		checkReturnStmt(scope, stmt)
		break
	case parser.ExprStmt:
		checkExprAllowVoid(scope, stmt.Expr)
		break
	}
}

func checkStmtBlock(scope *Scope, block parser.StmtBlock) {
	for _, stmt := range block.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkIfStmt(scope *Scope, stmt parser.IfStmt) {
	typ := checkExpr(scope, stmt.Cond)
	if typ.Equals(types.Bool) == false {
		scope.addError(fmt.Errorf("condition must resolve to a boolean"))
	}

	checkStmtBlock(scope, stmt.Clause)
}

func checkDeclarationStmt(scope *Scope, stmt parser.DeclarationStmt) {
	name := stmt.Name.Name
	typ := checkExpr(scope, stmt.Expr)
	scope.registerLocalVariable(name, typ)
}

func checkReturnStmt(scope *Scope, stmt parser.ReturnStmt) {
	var ret types.Type = types.TypeVoid{}
	if stmt.Expr != nil {
		ret = checkExpr(scope, stmt.Expr)
	}

	if scope.self == nil {
		scope.addError(fmt.Errorf("return statements must be inside a function"))
		return
	}

	if scope.self.Equals(ret) || ret.IsError() {
		return
	}

	if scope.self.Ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected to return nothing, got '%s'", ret))
		return
	}

	if ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected a return type of '%s', got nothing", scope.self.Ret))
	}

	scope.addError(fmt.Errorf("expected to return '%s', got '%s'", scope.self.Ret, ret))
}

func checkExprAllowVoid(scope *Scope, expr parser.Expr) types.Type {
	var typ types.Type = types.TypeError{}
	switch expr := expr.(type) {
	case parser.FunctionExpr:
		typ = checkFunctionExpr(scope, expr)
	case parser.DispatchExpr:
		typ = checkDispatchExpr(scope, expr)
	case parser.AssignExpr:
		typ = checkAssignExpr(scope, expr)
	case parser.BinaryExpr:
		typ = checkBinaryExpr(scope, expr)
	case parser.SubscriptExpr:
		typ = checkSubscriptExpr(scope, expr)
	case parser.SelfExpr:
		typ = checkSelfExpr(scope, expr)
	case parser.IdentExpr:
		typ = checkIdentExpr(scope, expr)
	case parser.NumberExpr:
		typ = checkNumberExpr(scope, expr)
	case parser.StringExpr:
		typ = checkStringExpr(scope, expr)
	case parser.BooleanExpr:
		typ = checkBooleanExpr(scope, expr)
	default:
		scope.addError(fmt.Errorf("unknown expression type"))
	}

	return typ
}

func checkExpr(scope *Scope, expr parser.Expr) types.Type {
	typ := checkExprAllowVoid(scope, expr)

	if typ.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("cannot use void types in an expression"))
		return types.TypeError{}
	}

	return typ
}

func checkFunctionExpr(scope *Scope, expr parser.FunctionExpr) types.Type {
	ret := types.ConvertTypeNote(expr.Ret)
	params := []types.Type{}
	for _, param := range expr.Params {
		params = append(params, types.ConvertTypeNote(param.Note))
	}
	tuple := types.TypeTuple{Children: params}
	self := types.TypeFunction{Params: tuple, Ret: ret}

	childScope := makeScope(scope, &self)

	for _, param := range expr.Params {
		paramName := param.Name.Name
		paramType := types.ConvertTypeNote(param.Note)
		childScope.registerLocalVariable(paramName, paramType)
	}

	checkStmtBlock(childScope, expr.Block)
	return self
}

func checkDispatchExpr(scope *Scope, expr parser.DispatchExpr) types.Type {
	// Resolve arguments to types
	argTypes := []types.Type{}
	for _, argExpr := range expr.Args {
		argTypes = append(argTypes, checkExpr(scope, argExpr))
	}

	// Resolve callee to type
	calleeType := checkExpr(scope, expr.Callee)
	calleeFunc, ok := calleeType.(types.TypeFunction)
	if ok == false {
		if calleeType.IsError() == false {
			scope.addError(fmt.Errorf("cannot call function on type '%s'", calleeType))
		}

		return types.TypeError{}
	}

	// Resolve return type
	retType := calleeFunc.Ret

	// Check that the given argument types match the expected parameter types
	totalArgs := len(argTypes)
	totalParams := len(calleeFunc.Params.Children)
	if totalArgs == totalParams {
		for i := 0; i < totalArgs; i++ {
			argType := argTypes[i]
			paramType := calleeFunc.Params.Children[i]

			if argType.IsError() {
				retType = types.TypeError{}
			} else if argType.Equals(paramType) == false {
				scope.addError(fmt.Errorf("expected '%s', got '%s'", paramType, argType))
				retType = types.TypeError{}
			}
		}
	} else {
		scope.addError(fmt.Errorf("expected %d arguments, got %d", totalParams, totalArgs))
		retType = types.TypeError{}
	}

	return retType
}

func checkAssignExpr(scope *Scope, expr parser.AssignExpr) types.Type {
	name := expr.Left.Name
	leftType := scope.getVariable(name)
	rightType := checkExpr(scope, expr.Right)

	if leftType == nil {
		scope.addError(fmt.Errorf("'%s' cannot be assigned before it is declared", name))
		return types.TypeError{}
	}

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if leftType.Equals(rightType) == false {
		scope.addError(fmt.Errorf("'%s' cannot be assigned type '%s'", leftType, rightType))
		return types.TypeError{}
	}

	return leftType
}

func checkBinaryExpr(scope *Scope, expr parser.BinaryExpr) types.Type {
	leftType := checkExpr(scope, expr.Left)
	rightType := checkExpr(scope, expr.Right)

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if operLUT, ok := binops[expr.Oper]; ok {
		if leftLUT, ok := operLUT[leftType]; ok {
			if retType, ok := leftLUT[rightType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("operator '%s' does not support %s and %s", expr.Oper, leftType, rightType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '%s'", expr.Oper))
	return types.TypeError{}
}

func checkSubscriptExpr(scope *Scope, expr parser.SubscriptExpr) types.Type {
	listType := checkExpr(scope, expr.ListLike)
	indexType := checkExpr(scope, expr.Index)

	if listType.IsError() || indexType.IsError() {
		return types.TypeError{}
	}

	if subscriptLUT, ok := binops["["]; ok {
		if listLUT, ok := subscriptLUT[listType]; ok {
			if retType, ok := listLUT[indexType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("subscript operator does not support %s[%s]", listType, indexType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '['"))
	return types.TypeError{}
}

func checkSelfExpr(scope *Scope, expr parser.SelfExpr) types.Type {
	if scope.self == nil {
		scope.addError(fmt.Errorf("self references must be inside a function"))
		return types.TypeError{}
	}

	return *scope.self
}

func checkIdentExpr(scope *Scope, expr parser.IdentExpr) types.Type {
	if scope.existingVariable(expr.Name) {
		return scope.getVariable(expr.Name)
	}

	scope.addError(fmt.Errorf("variable '%s' was used before it was declared", expr.Name))
	return types.TypeError{}
}

func checkNumberExpr(scope *Scope, expr parser.NumberExpr) types.Type {
	return types.Int
}

func checkStringExpr(scope *Scope, expr parser.StringExpr) types.Type {
	return types.Str
}

func checkBooleanExpr(scope *Scope, expr parser.BooleanExpr) types.Type {
	return types.Bool
}
<commit_msg>use type aliases for binary operation lookup-table<commit_after>package check

import (
	"fmt"
	"plaid/parser"
	"plaid/types"
	"plaid/vm"
)

type binopsLUT map[string]map[types.Type]map[types.Type]types.Type
type doubleLUT map[types.Type]map[types.Type]types.Type
type singleLUT map[types.Type]types.Type

var binops = binopsLUT{
	"+": doubleLUT{
		types.Int: singleLUT{types.Int: types.Int},
		types.Str: singleLUT{types.Str: types.Str},
	},
	"-": doubleLUT{
		types.Int: singleLUT{types.Int: types.Int},
	},
	"*": doubleLUT{
		types.Int: singleLUT{types.Int: types.Int},
	},
	"<": doubleLUT{
		types.Int: singleLUT{types.Int: types.Bool},
	},
	"<=": doubleLUT{
		types.Int: singleLUT{types.Int: types.Bool},
	},
	">": doubleLUT{
		types.Int: singleLUT{types.Int: types.Bool},
	},
	">=": doubleLUT{
		types.Int: singleLUT{types.Int: types.Bool},
	},
	"[": doubleLUT{
		types.Str: singleLUT{types.Int: types.Str},
	},
}

// Check takes an existing abstract syntax tree and performs type checks and
// other correctness checks. It returns a list of any errors that were
// discovered inside the AST
func Check(prog parser.Program, libraries ...vm.Library) *Scope {
	global := makeScope(nil, nil)

	for _, library := range libraries {
		for name, builtin := range library {
			global.registerLocalVariable(name, builtin.Type)
		}
	}

	checkProgram(global, prog)
	return global
}

func checkProgram(scope *Scope, prog parser.Program) {
	for _, stmt := range prog.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkStmt(scope *Scope, stmt parser.Stmt) {
	switch stmt := stmt.(type) {
	case parser.IfStmt:
		checkIfStmt(scope, stmt)
		break
	case parser.DeclarationStmt:
		checkDeclarationStmt(scope, stmt)
		break
	case parser.ReturnStmt:
		checkReturnStmt(scope, stmt)
		break
	case parser.ExprStmt:
		checkExprAllowVoid(scope, stmt.Expr)
		break
	}
}

func checkStmtBlock(scope *Scope, block parser.StmtBlock) {
	for _, stmt := range block.Stmts {
		checkStmt(scope, stmt)
	}
}

func checkIfStmt(scope *Scope, stmt parser.IfStmt) {
	typ := checkExpr(scope, stmt.Cond)
	if typ.Equals(types.Bool) == false {
		scope.addError(fmt.Errorf("condition must resolve to a boolean"))
	}

	checkStmtBlock(scope, stmt.Clause)
}

func checkDeclarationStmt(scope *Scope, stmt parser.DeclarationStmt) {
	name := stmt.Name.Name
	typ := checkExpr(scope, stmt.Expr)
	scope.registerLocalVariable(name, typ)
}

func checkReturnStmt(scope *Scope, stmt parser.ReturnStmt) {
	var ret types.Type = types.TypeVoid{}
	if stmt.Expr != nil {
		ret = checkExpr(scope, stmt.Expr)
	}

	if scope.self == nil {
		scope.addError(fmt.Errorf("return statements must be inside a function"))
		return
	}

	if scope.self.Equals(ret) || ret.IsError() {
		return
	}

	if scope.self.Ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected to return nothing, got '%s'", ret))
		return
	}

	if ret.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("expected a return type of '%s', got nothing", scope.self.Ret))
	}

	scope.addError(fmt.Errorf("expected to return '%s', got '%s'", scope.self.Ret, ret))
}

func checkExprAllowVoid(scope *Scope, expr parser.Expr) types.Type {
	var typ types.Type = types.TypeError{}
	switch expr := expr.(type) {
	case parser.FunctionExpr:
		typ = checkFunctionExpr(scope, expr)
	case parser.DispatchExpr:
		typ = checkDispatchExpr(scope, expr)
	case parser.AssignExpr:
		typ = checkAssignExpr(scope, expr)
	case parser.BinaryExpr:
		typ = checkBinaryExpr(scope, expr)
	case parser.SubscriptExpr:
		typ = checkSubscriptExpr(scope, expr)
	case parser.SelfExpr:
		typ = checkSelfExpr(scope, expr)
	case parser.IdentExpr:
		typ = checkIdentExpr(scope, expr)
	case parser.NumberExpr:
		typ = checkNumberExpr(scope, expr)
	case parser.StringExpr:
		typ = checkStringExpr(scope, expr)
	case parser.BooleanExpr:
		typ = checkBooleanExpr(scope, expr)
	default:
		scope.addError(fmt.Errorf("unknown expression type"))
	}

	return typ
}

func checkExpr(scope *Scope, expr parser.Expr) types.Type {
	typ := checkExprAllowVoid(scope, expr)

	if typ.Equals(types.TypeVoid{}) {
		scope.addError(fmt.Errorf("cannot use void types in an expression"))
		return types.TypeError{}
	}

	return typ
}

func checkFunctionExpr(scope *Scope, expr parser.FunctionExpr) types.Type {
	ret := types.ConvertTypeNote(expr.Ret)
	params := []types.Type{}
	for _, param := range expr.Params {
		params = append(params, types.ConvertTypeNote(param.Note))
	}
	tuple := types.TypeTuple{Children: params}
	self := types.TypeFunction{Params: tuple, Ret: ret}

	childScope := makeScope(scope, &self)

	for _, param := range expr.Params {
		paramName := param.Name.Name
		paramType := types.ConvertTypeNote(param.Note)
		childScope.registerLocalVariable(paramName, paramType)
	}

	checkStmtBlock(childScope, expr.Block)
	return self
}

func checkDispatchExpr(scope *Scope, expr parser.DispatchExpr) types.Type {
	// Resolve arguments to types
	argTypes := []types.Type{}
	for _, argExpr := range expr.Args {
		argTypes = append(argTypes, checkExpr(scope, argExpr))
	}

	// Resolve callee to type
	calleeType := checkExpr(scope, expr.Callee)
	calleeFunc, ok := calleeType.(types.TypeFunction)
	if ok == false {
		if calleeType.IsError() == false {
			scope.addError(fmt.Errorf("cannot call function on type '%s'", calleeType))
		}

		return types.TypeError{}
	}

	// Resolve return type
	retType := calleeFunc.Ret

	// Check that the given argument types match the expected parameter types
	totalArgs := len(argTypes)
	totalParams := len(calleeFunc.Params.Children)
	if totalArgs == totalParams {
		for i := 0; i < totalArgs; i++ {
			argType := argTypes[i]
			paramType := calleeFunc.Params.Children[i]

			if argType.IsError() {
				retType = types.TypeError{}
			} else if argType.Equals(paramType) == false {
				scope.addError(fmt.Errorf("expected '%s', got '%s'", paramType, argType))
				retType = types.TypeError{}
			}
		}
	} else {
		scope.addError(fmt.Errorf("expected %d arguments, got %d", totalParams, totalArgs))
		retType = types.TypeError{}
	}

	return retType
}

func checkAssignExpr(scope *Scope, expr parser.AssignExpr) types.Type {
	name := expr.Left.Name
	leftType := scope.getVariable(name)
	rightType := checkExpr(scope, expr.Right)

	if leftType == nil {
		scope.addError(fmt.Errorf("'%s' cannot be assigned before it is declared", name))
		return types.TypeError{}
	}

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if leftType.Equals(rightType) == false {
		scope.addError(fmt.Errorf("'%s' cannot be assigned type '%s'", leftType, rightType))
		return types.TypeError{}
	}

	return leftType
}

func checkBinaryExpr(scope *Scope, expr parser.BinaryExpr) types.Type {
	leftType := checkExpr(scope, expr.Left)
	rightType := checkExpr(scope, expr.Right)

	if leftType.IsError() || rightType.IsError() {
		return types.TypeError{}
	}

	if operLUT, ok := binops[expr.Oper]; ok {
		if leftLUT, ok := operLUT[leftType]; ok {
			if retType, ok := leftLUT[rightType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("operator '%s' does not support %s and %s", expr.Oper, leftType, rightType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '%s'", expr.Oper))
	return types.TypeError{}
}

func checkSubscriptExpr(scope *Scope, expr parser.SubscriptExpr) types.Type {
	listType := checkExpr(scope, expr.ListLike)
	indexType := checkExpr(scope, expr.Index)

	if listType.IsError() || indexType.IsError() {
		return types.TypeError{}
	}

	if subscriptLUT, ok := binops["["]; ok {
		if listLUT, ok := subscriptLUT[listType]; ok {
			if retType, ok := listLUT[indexType]; ok {
				return retType
			}
		}

		scope.addError(fmt.Errorf("subscript operator does not support %s[%s]", listType, indexType))
		return types.TypeError{}
	}

	scope.addError(fmt.Errorf("unknown infix operator '['"))
	return types.TypeError{}
}

func checkSelfExpr(scope *Scope, expr parser.SelfExpr) types.Type {
	if scope.self == nil {
		scope.addError(fmt.Errorf("self references must be inside a function"))
		return types.TypeError{}
	}

	return *scope.self
}

func checkIdentExpr(scope *Scope, expr parser.IdentExpr) types.Type {
	if scope.existingVariable(expr.Name) {
		return scope.getVariable(expr.Name)
	}

	scope.addError(fmt.Errorf("variable '%s' was used before it was declared", expr.Name))
	return types.TypeError{}
}

func checkNumberExpr(scope *Scope, expr parser.NumberExpr) types.Type {
	return types.Int
}

func checkStringExpr(scope *Scope, expr parser.StringExpr) types.Type {
	return types.Str
}

func checkBooleanExpr(scope *Scope, expr parser.BooleanExpr) types.Type {
	return types.Bool
}
<|endoftext|>


<commit_before>package dockertest_test

import (
	"database/sql"
	"fmt"
	"net/http"
	"strings"
	"testing"
	"time"

	"gopkg.in/mgo.v2"

	"github.com/garyburd/redigo/redis"
	. "github.com/ninnemana/dockertest"
	"github.com/ninnemana/elastigo/lib"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestOpenPostgreSQLContainerConnection(t *testing.T) {
	c, db, err := OpenPostgreSQLContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.Nil(t, db.Ping())
	require.NotNil(t, db)
	defer db.Close()
}

func TestOpenMySQLContainerConnection(t *testing.T) {
	c, db, err := OpenMySQLContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.Nil(t, db.Ping())
	require.NotNil(t, db)
	defer db.Close()
}

func TestOpenMongoDBContainerConnection(t *testing.T) {
	c, db, err := OpenMongoDBContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, db)
	_, err = db.DatabaseNames()
	require.Nil(t, err)
	defer db.Close()
}

func TestOpenElasticSearchContainerConnection(t *testing.T) {
	c, conn, err := OpenElasticSearchContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, conn)
	_, err = conn.Health("")
	require.Nil(t, err)
	defer conn.Close()
}

func TestOpenRedisContainerConnection(t *testing.T) {
	c, client, err := OpenRedisContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, client)

	v, err := client.Cmd("echo", "Hello, World!").Str()
	require.Nil(t, err)
	assert.Equal(t, "Hello, World!", v)

	defer client.Close()
}

func TestOpenNSQLookupdContainerConnection(t *testing.T) {

	c, ip, tcpPort, httpPort, err := OpenNSQLookupdContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotEmpty(t, ip)
	require.NotZero(t, tcpPort)
	require.NotZero(t, httpPort)

	resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
	require.Nil(t, err)
	require.Equal(t, resp.StatusCode, 200)

}

func TestOpenNSQContainerConnection(t *testing.T) {

	c, ip, tcpPort, httpPort, err := OpenNSQdContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotEmpty(t, ip)
	require.NotZero(t, tcpPort)
	require.NotZero(t, httpPort)

	resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
	require.Nil(t, err)
	require.Equal(t, resp.StatusCode, 200)

}

func TestConnectToPostgreSQL(t *testing.T) {
	c, err := ConnectToPostgreSQL(15, time.Millisecond*500, func(url string) bool {
		db, err := sql.Open("postgres", url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToMySQL(t *testing.T) {
	c, err := ConnectToMySQL(15, time.Millisecond*500, func(url string) bool {
		db, err := sql.Open("mysql", url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToMongoDB(t *testing.T) {
	c, err := ConnectToMongoDB(15, time.Millisecond*500, func(url string) bool {
		db, err := mgo.Dial(url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToElasticSearch(t *testing.T) {
	c, err := ConnectToElasticSearch(15, time.Millisecond*500, func(url string) bool {
		segs := strings.Split(url, ":")
		if len(segs) != 2 {
			return false
		}

		conn := elastigo.NewConn()
		conn.Domain = segs[0]
		conn.Port = segs[1]
		resp, err := conn.Health()
		if err != nil {
			return false
		}
		if resp.Status != "green" {
			return false
		}
		// defer conn.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToRedis(t *testing.T) {
	c, err := ConnectToRedis(15, time.Millisecond*500, func(url string) bool {
		client, err := redis.DialTimeout("tcp", url, 10*time.Second, 10*time.Second, 10*time.Second)
		require.Nil(t, err)
		require.NotNil(t, client)

		reply, err := redis.String(client.Do("echo", "Hello, World!"))

		require.Nil(t, err)
		assert.Equal(t, "Hello, World!", reply)

		defer client.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToNSQLookupd(t *testing.T) {
	c, err := ConnectToNSQLookupd(15, time.Millisecond*500, func(ip string, httpPort int, tcpPort int) bool {
		resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
		require.Nil(t, err)
		require.Equal(t, resp.StatusCode, 200)

		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToNSQd(t *testing.T) {
	c, err := ConnectToNSQd(15, time.Millisecond*500, func(ip string, httpPort int, tcpPort int) bool {
		resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
		require.Nil(t, err)
		require.Equal(t, resp.StatusCode, 200)

		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}
<commit_msg>changed import back<commit_after>package dockertest_test

import (
	"database/sql"
	"fmt"
	"net/http"
	"strings"
	"testing"
	"time"

	"gopkg.in/mgo.v2"

	"github.com/garyburd/redigo/redis"
	. "github.com/ninnemana/dockertest"
	"github.com/ory-am/elastigo/lib"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestOpenPostgreSQLContainerConnection(t *testing.T) {
	c, db, err := OpenPostgreSQLContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.Nil(t, db.Ping())
	require.NotNil(t, db)
	defer db.Close()
}

func TestOpenMySQLContainerConnection(t *testing.T) {
	c, db, err := OpenMySQLContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.Nil(t, db.Ping())
	require.NotNil(t, db)
	defer db.Close()
}

func TestOpenMongoDBContainerConnection(t *testing.T) {
	c, db, err := OpenMongoDBContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, db)
	_, err = db.DatabaseNames()
	require.Nil(t, err)
	defer db.Close()
}

func TestOpenElasticSearchContainerConnection(t *testing.T) {
	c, conn, err := OpenElasticSearchContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, conn)
	_, err = conn.Health("")
	require.Nil(t, err)
	defer conn.Close()
}

func TestOpenRedisContainerConnection(t *testing.T) {
	c, client, err := OpenRedisContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotNil(t, client)

	v, err := client.Cmd("echo", "Hello, World!").Str()
	require.Nil(t, err)
	assert.Equal(t, "Hello, World!", v)

	defer client.Close()
}

func TestOpenNSQLookupdContainerConnection(t *testing.T) {

	c, ip, tcpPort, httpPort, err := OpenNSQLookupdContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotEmpty(t, ip)
	require.NotZero(t, tcpPort)
	require.NotZero(t, httpPort)

	resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
	require.Nil(t, err)
	require.Equal(t, resp.StatusCode, 200)

}

func TestOpenNSQContainerConnection(t *testing.T) {

	c, ip, tcpPort, httpPort, err := OpenNSQdContainerConnection(15, time.Millisecond*500)
	require.Nil(t, err)
	defer c.KillRemove()
	require.NotEmpty(t, ip)
	require.NotZero(t, tcpPort)
	require.NotZero(t, httpPort)

	resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
	require.Nil(t, err)
	require.Equal(t, resp.StatusCode, 200)

}

func TestConnectToPostgreSQL(t *testing.T) {
	c, err := ConnectToPostgreSQL(15, time.Millisecond*500, func(url string) bool {
		db, err := sql.Open("postgres", url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToMySQL(t *testing.T) {
	c, err := ConnectToMySQL(15, time.Millisecond*500, func(url string) bool {
		db, err := sql.Open("mysql", url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToMongoDB(t *testing.T) {
	c, err := ConnectToMongoDB(15, time.Millisecond*500, func(url string) bool {
		db, err := mgo.Dial(url)
		if err != nil {
			return false
		}
		defer db.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToElasticSearch(t *testing.T) {
	c, err := ConnectToElasticSearch(15, time.Millisecond*500, func(url string) bool {
		segs := strings.Split(url, ":")
		if len(segs) != 2 {
			return false
		}

		conn := elastigo.NewConn()
		conn.Domain = segs[0]
		conn.Port = segs[1]
		resp, err := conn.Health()
		if err != nil {
			return false
		}
		if resp.Status != "green" {
			return false
		}
		// defer conn.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToRedis(t *testing.T) {
	c, err := ConnectToRedis(15, time.Millisecond*500, func(url string) bool {
		client, err := redis.DialTimeout("tcp", url, 10*time.Second, 10*time.Second, 10*time.Second)
		require.Nil(t, err)
		require.NotNil(t, client)

		reply, err := redis.String(client.Do("echo", "Hello, World!"))

		require.Nil(t, err)
		assert.Equal(t, "Hello, World!", reply)

		defer client.Close()
		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToNSQLookupd(t *testing.T) {
	c, err := ConnectToNSQLookupd(15, time.Millisecond*500, func(ip string, httpPort int, tcpPort int) bool {
		resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
		require.Nil(t, err)
		require.Equal(t, resp.StatusCode, 200)

		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}

func TestConnectToNSQd(t *testing.T) {
	c, err := ConnectToNSQd(15, time.Millisecond*500, func(ip string, httpPort int, tcpPort int) bool {
		resp, err := http.Get(fmt.Sprintf("http://%s:%d/ping", ip, httpPort))
		require.Nil(t, err)
		require.Equal(t, resp.StatusCode, 200)

		return true
	})
	require.Nil(t, err)
	defer c.KillRemove()
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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 main

import (
	"fmt"
	"github.com/golang/oauth2"
	"net/http"
	"os"
)

type gcloudTransport struct {
	transport http.RoundTripper
}

func (gt gcloudTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	req.Header.Add("User-Agent", "Cloud SDK Command Line Tool")
	req.Header.Add("User-Agent", "gcloud")
	req.Header.Add("User-Agent", "gcloud_apis")
	return gt.transport.RoundTrip(req)
}

/*
This auth function is a bootstrap until better integration with
gcloud auth can be provided. To use, set (and export) the environment
variable GCLOUD_APIS_REFRESH_TOKEN to a valid refresh token, which can
usually be acquired via $(gcloud auth print-refresh-token). In the
future, gcloud_apis will be able to fetch credential information from
the same place as gcloud, and will be usable without setting this
environment variable.
*/
func getAuthenticatedClient() (*http.Client, error) {
	flow, err := oauth2.New(
		oauth2.Client(
			"32555940559.apps.googleusercontent.com",
			"ZmssLNjJy2998hD4CTg2ejr2",
		),
		oauth2.Scope(
			"https://www.googleapis.com/auth/cloud-platform",
		),
		oauth2.Endpoint(
			"https://accounts.google.com/o/oauth2/auth",
			"https://accounts.google.com/o/oauth2/token",
		),
	)
	if err != nil {
		return nil, err
	}

	refreshToken := os.Getenv("GCLOUD_APIS_REFRESH_TOKEN")
	if refreshToken == "" {
		err = fmt.Errorf("no refresh token found in GCLOUD_APIS_REFRESH_TOKEN")
		return nil, err
	}

	token := &oauth2.Token{
		RefreshToken: refreshToken,
	}

	oauth2Transport := flow.NewTransportFromToken(token)

	client := &http.Client{
		Transport: gcloudTransport{oauth2Transport},
	}

	return client, nil
}
<commit_msg>Updated auth.go to match updates in oauth2 library.<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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 main

import (
	"fmt"
	"github.com/golang/oauth2"
	"net/http"
	"os"
)

type memoryTokenStore struct {
	token *oauth2.Token
}

func (s memoryTokenStore) ReadToken() (*oauth2.Token, error) {
	return s.token, nil
}

func (s memoryTokenStore) WriteToken(*oauth2.Token) {
	return
}

type gcloudTransport struct {
	transport http.RoundTripper
}

func (gt gcloudTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	req.Header.Add("User-Agent", "Cloud SDK Command Line Tool")
	req.Header.Add("User-Agent", "gcloud")
	req.Header.Add("User-Agent", "gcloud_apis")
	return gt.transport.RoundTrip(req)
}

/*
This auth function is a bootstrap until better integration with
gcloud auth can be provided. To use, set (and export) the environment
variable GCLOUD_APIS_REFRESH_TOKEN to a valid refresh token, which can
usually be acquired via $(gcloud auth print-refresh-token). In the
future, gcloud_apis will be able to fetch credential information from
the same place as gcloud, and will be usable without setting this
environment variable.
*/
func getAuthenticatedClient() (*http.Client, error) {
	flow, err := oauth2.New(
		oauth2.Client(
			"32555940559.apps.googleusercontent.com",
			"ZmssLNjJy2998hD4CTg2ejr2",
		),
		oauth2.Scope(
			"https://www.googleapis.com/auth/cloud-platform",
		),
		oauth2.Endpoint(
			"https://accounts.google.com/o/oauth2/auth",
			"https://accounts.google.com/o/oauth2/token",
		),
	)
	if err != nil {
		return nil, err
	}

	refreshToken := os.Getenv("GCLOUD_APIS_REFRESH_TOKEN")
	if refreshToken == "" {
		err = fmt.Errorf("no refresh token found in GCLOUD_APIS_REFRESH_TOKEN")
		return nil, err
	}

	token := &oauth2.Token{
		RefreshToken: refreshToken,
	}

	oauth2Transport, err := flow.NewTransportFromTokenStore(memoryTokenStore{token})
	if err != nil {
		return nil, err
	}

	client := &http.Client{
		Transport: gcloudTransport{oauth2Transport},
	}

	return client, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Space Monkey, 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.

// +build cgo,!darwin

package openssl

// #include <openssl/evp.h>
import "C"

import (
	"errors"
	"fmt"
)

type AuthenticatedEncryptionCipherCtx interface {
	EncryptionCipherCtx

	// data passed in to ExtraData() is part of the final output; it is
	// not encrypted itself, but is part of the authenticated data. when
	// decrypting or authenticating, pass back with the decryption
	// context's ExtraData()
	ExtraData([]byte) error

	// use after finalizing encryption to get the authenticating tag
	GetTag() ([]byte, error)
}

type AuthenticatedDecryptionCipherCtx interface {
	DecryptionCipherCtx

	// pass in any extra data that was added during encryption with the
	// encryption context's ExtraData()
	ExtraData([]byte) error

	// use before finalizing decryption to tell the library what the
	// tag is expected to be
	SetTag([]byte) error
}

type authEncryptionCipherCtx struct {
	*encryptionCipherCtx
}

type authDecryptionCipherCtx struct {
	*decryptionCipherCtx
}

func getGCMCipher(blocksize int) (*Cipher, error) {
	var cipherptr *C.EVP_CIPHER
	switch blocksize {
	case 256:
		cipherptr = C.EVP_aes_256_gcm()
	case 192:
		cipherptr = C.EVP_aes_192_gcm()
	case 128:
		cipherptr = C.EVP_aes_128_gcm()
	default:
		return nil, fmt.Errorf("unknown block size %d", blocksize)
	}
	return &Cipher{ptr: cipherptr}, nil
}

func NewGCMEncryptionCipherCtx(blocksize int, e *Engine, key, iv []byte) (
	AuthenticatedEncryptionCipherCtx, error) {
	cipher, err := getGCMCipher(blocksize)
	if err != nil {
		return nil, err
	}
	ctx, err := newEncryptionCipherCtx(cipher, e, key, nil)
	if err != nil {
		return nil, err
	}
	if iv != nil {
		err := ctx.setCtrl(C.EVP_CTRL_GCM_SET_IVLEN, len(iv))
		if err != nil {
			return nil, fmt.Errorf("could not set IV len to %d: %s",
				len(iv), err)
		}
		if 1 != C.EVP_EncryptInit_ex(ctx.ctx, nil, nil, nil,
			(*C.uchar)(&iv[0])) {
			return nil, errors.New("failed to apply IV")
		}
	}
	return &authEncryptionCipherCtx{encryptionCipherCtx: ctx}, nil
}

func NewGCMDecryptionCipherCtx(blocksize int, e *Engine, key, iv []byte) (
	AuthenticatedDecryptionCipherCtx, error) {
	cipher, err := getGCMCipher(blocksize)
	if err != nil {
		return nil, err
	}
	ctx, err := newDecryptionCipherCtx(cipher, e, key, nil)
	if err != nil {
		return nil, err
	}
	if iv != nil {
		err := ctx.setCtrl(C.EVP_CTRL_GCM_SET_IVLEN, len(iv))
		if err != nil {
			return nil, fmt.Errorf("could not set IV len to %d: %s",
				len(iv), err)
		}
		if 1 != C.EVP_DecryptInit_ex(ctx.ctx, nil, nil, nil,
			(*C.uchar)(&iv[0])) {
			return nil, errors.New("failed to apply IV")
		}
	}
	return &authDecryptionCipherCtx{decryptionCipherCtx: ctx}, nil
}

func (ctx *authEncryptionCipherCtx) ExtraData(aad []byte) error {
	if aad == nil {
		return nil
	}
	var outlen C.int
	if 1 != C.EVP_EncryptUpdate(ctx.ctx, nil, &outlen, (*C.uchar)(&aad[0]),
		C.int(len(aad))) {
		return errors.New("failed to add additional authenticated data")
	}
	return nil
}

func (ctx *authDecryptionCipherCtx) ExtraData(aad []byte) error {
	if aad == nil {
		return nil
	}
	var outlen C.int
	if 1 != C.EVP_DecryptUpdate(ctx.ctx, nil, &outlen, (*C.uchar)(&aad[0]),
		C.int(len(aad))) {
		return errors.New("failed to add additional authenticated data")
	}
	return nil
}

func (ctx *authEncryptionCipherCtx) GetTag() ([]byte, error) {
	return ctx.getCtrlBytes(C.EVP_CTRL_GCM_GET_TAG, GCM_TAG_MAXLEN,
		GCM_TAG_MAXLEN)
}

func (ctx *authDecryptionCipherCtx) SetTag(tag []byte) error {
	return ctx.setCtrlBytes(C.EVP_CTRL_GCM_SET_TAG, len(tag), tag)
}
<commit_msg>fix taking address of first element of empty byte slice<commit_after>// Copyright (C) 2014 Space Monkey, 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.

// +build cgo,!darwin

package openssl

// #include <openssl/evp.h>
import "C"

import (
	"errors"
	"fmt"
)

type AuthenticatedEncryptionCipherCtx interface {
	EncryptionCipherCtx

	// data passed in to ExtraData() is part of the final output; it is
	// not encrypted itself, but is part of the authenticated data. when
	// decrypting or authenticating, pass back with the decryption
	// context's ExtraData()
	ExtraData([]byte) error

	// use after finalizing encryption to get the authenticating tag
	GetTag() ([]byte, error)
}

type AuthenticatedDecryptionCipherCtx interface {
	DecryptionCipherCtx

	// pass in any extra data that was added during encryption with the
	// encryption context's ExtraData()
	ExtraData([]byte) error

	// use before finalizing decryption to tell the library what the
	// tag is expected to be
	SetTag([]byte) error
}

type authEncryptionCipherCtx struct {
	*encryptionCipherCtx
}

type authDecryptionCipherCtx struct {
	*decryptionCipherCtx
}

func getGCMCipher(blocksize int) (*Cipher, error) {
	var cipherptr *C.EVP_CIPHER
	switch blocksize {
	case 256:
		cipherptr = C.EVP_aes_256_gcm()
	case 192:
		cipherptr = C.EVP_aes_192_gcm()
	case 128:
		cipherptr = C.EVP_aes_128_gcm()
	default:
		return nil, fmt.Errorf("unknown block size %d", blocksize)
	}
	return &Cipher{ptr: cipherptr}, nil
}

func NewGCMEncryptionCipherCtx(blocksize int, e *Engine, key, iv []byte) (
	AuthenticatedEncryptionCipherCtx, error) {
	cipher, err := getGCMCipher(blocksize)
	if err != nil {
		return nil, err
	}
	ctx, err := newEncryptionCipherCtx(cipher, e, key, nil)
	if err != nil {
		return nil, err
	}
	if len(iv) > 0 {
		err := ctx.setCtrl(C.EVP_CTRL_GCM_SET_IVLEN, len(iv))
		if err != nil {
			return nil, fmt.Errorf("could not set IV len to %d: %s",
				len(iv), err)
		}
		if 1 != C.EVP_EncryptInit_ex(ctx.ctx, nil, nil, nil,
			(*C.uchar)(&iv[0])) {
			return nil, errors.New("failed to apply IV")
		}
	}
	return &authEncryptionCipherCtx{encryptionCipherCtx: ctx}, nil
}

func NewGCMDecryptionCipherCtx(blocksize int, e *Engine, key, iv []byte) (
	AuthenticatedDecryptionCipherCtx, error) {
	cipher, err := getGCMCipher(blocksize)
	if err != nil {
		return nil, err
	}
	ctx, err := newDecryptionCipherCtx(cipher, e, key, nil)
	if err != nil {
		return nil, err
	}
	if len(iv) > 0 {
		err := ctx.setCtrl(C.EVP_CTRL_GCM_SET_IVLEN, len(iv))
		if err != nil {
			return nil, fmt.Errorf("could not set IV len to %d: %s",
				len(iv), err)
		}
		if 1 != C.EVP_DecryptInit_ex(ctx.ctx, nil, nil, nil,
			(*C.uchar)(&iv[0])) {
			return nil, errors.New("failed to apply IV")
		}
	}
	return &authDecryptionCipherCtx{decryptionCipherCtx: ctx}, nil
}

func (ctx *authEncryptionCipherCtx) ExtraData(aad []byte) error {
	if aad == nil {
		return nil
	}
	var outlen C.int
	if 1 != C.EVP_EncryptUpdate(ctx.ctx, nil, &outlen, (*C.uchar)(&aad[0]),
		C.int(len(aad))) {
		return errors.New("failed to add additional authenticated data")
	}
	return nil
}

func (ctx *authDecryptionCipherCtx) ExtraData(aad []byte) error {
	if aad == nil {
		return nil
	}
	var outlen C.int
	if 1 != C.EVP_DecryptUpdate(ctx.ctx, nil, &outlen, (*C.uchar)(&aad[0]),
		C.int(len(aad))) {
		return errors.New("failed to add additional authenticated data")
	}
	return nil
}

func (ctx *authEncryptionCipherCtx) GetTag() ([]byte, error) {
	return ctx.getCtrlBytes(C.EVP_CTRL_GCM_GET_TAG, GCM_TAG_MAXLEN,
		GCM_TAG_MAXLEN)
}

func (ctx *authDecryptionCipherCtx) SetTag(tag []byte) error {
	return ctx.setCtrlBytes(C.EVP_CTRL_GCM_SET_TAG, len(tag), tag)
}
<|endoftext|>