<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 tcontainer

import (
	"reflect"
	"runtime"
	"sync/atomic"
	"unsafe"
)

const (
	tiny   = 64
	small  = 512
	medium = 1024
	large  = 1024 * 10
	huge   = 1024 * 100

	tinyCount   = 16384 // 1 MB
	smallCount  = 2048  // 1 MB
	mediumCount = 1024  // 1 MB
	largeCount  = 102   // ~1 MB
	hugeCount   = 10    // ~1 MB
)

type byteSlab struct {
	buffer     []byte
	bufferSize uintptr
	stride     uintptr
	basePtr    *uintptr
	nextPtr    *uintptr
}

// BytePool is a fragmentation friendly way to allocated byte slices.
type BytePool struct {
	tinySlab   byteSlab
	smallSlab  byteSlab
	mediumSlab byteSlab
	largeSlab  byteSlab
	hugeSlab   byteSlab
}

func newByteSlab(size, count int) byteSlab {
	bufferSize := count * size
	buffer := make([]byte, bufferSize)
	basePtr := (*reflect.SliceHeader)(unsafe.Pointer(&buffer)).Data
	nextPtr := basePtr + uintptr(bufferSize)

	return byteSlab{
		buffer:     buffer,
		bufferSize: uintptr(bufferSize),
		stride:     uintptr(size),
		basePtr:    &basePtr,
		nextPtr:    &nextPtr,
	}
}

func (slab *byteSlab) getSlice(size int) (chunk []byte) {
	chunkHeader := (*reflect.SliceHeader)(unsafe.Pointer(&chunk))
	chunkHeader.Len = size
	chunkHeader.Cap = int(slab.stride)

	for {
		// WARNING: The following two lines are order sensitive
		basePtr := atomic.LoadUintptr(slab.basePtr)
		nextPtr := atomic.AddUintptr(slab.nextPtr, -slab.stride)
		lastPtr := basePtr + slab.bufferSize

		switch {
		case nextPtr < basePtr || nextPtr >= lastPtr:
			// out of range either means alloc while realloc or race between
			// base and next during realloc. In the latter case we lose a chunk.
			runtime.Gosched()

		case nextPtr == basePtr:
			// Last item: realloc
			buffer := make([]byte, slab.bufferSize)
			dataPtr := (*reflect.SliceHeader)(unsafe.Pointer(&buffer)).Data

			// WARNING: The following two lines are order sensitive
			atomic.StoreUintptr(slab.nextPtr, dataPtr+slab.bufferSize)
			atomic.StoreUintptr(slab.basePtr, dataPtr)

			slab.buffer = buffer // hold the buffer so that GC does not collect
			fallthrough

		default:
			chunkHeader.Data = nextPtr
			return
		}
	}
}

// NewBytePool creates a new BytePool with each slab using 1 MB of storage.
// The pool contains 5 slabs of different sizes: 64B, 512B, 1KB, 10KB and 100KB.
// Allocations above 100KB will be allocated directly.
func NewBytePool() BytePool {
	return BytePool{
		tinySlab:   newByteSlab(tiny, tinyCount),
		smallSlab:  newByteSlab(small, smallCount),
		mediumSlab: newByteSlab(medium, mediumCount),
		largeSlab:  newByteSlab(large, largeCount),
		hugeSlab:   newByteSlab(huge, hugeCount),
	}
}

// NewBytePoolWithSize creates a new BytePool with each slab size using n MB of
// storage. See NewBytePool() for slab size details.
func NewBytePoolWithSize(n int) BytePool {
	if n <= 0 {
		n = 1
	}
	return BytePool{
		tinySlab:   newByteSlab(tiny, tinyCount*n),
		smallSlab:  newByteSlab(small, smallCount*n),
		mediumSlab: newByteSlab(medium, mediumCount*n),
		largeSlab:  newByteSlab(large, largeCount*n),
		hugeSlab:   newByteSlab(huge, hugeCount*n),
	}
}

// Get returns a slice allocated to a normalized size.
// Sizes are organized in evenly sized buckets so that fragmentation is kept low.
func (b *BytePool) Get(size int) []byte {
	switch {
	case size == 0:
		return []byte{}

	case size <= tiny:
		return b.tinySlab.getSlice(size)

	case size <= small:
		return b.smallSlab.getSlice(size)

	case size <= medium:
		return b.mediumSlab.getSlice(size)

	case size <= large:
		return b.largeSlab.getSlice(size)

	case size <= huge:
		return b.hugeSlab.getSlice(size)

	default:
		return make([]byte, size)
	}
}
<commit_msg>fixed a race condition when overwriting the byte buffer hook<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 tcontainer

import (
	"reflect"
	"runtime"
	"sync/atomic"
	"unsafe"
)

const (
	tiny   = 64
	small  = 512
	medium = 1024
	large  = 1024 * 10
	huge   = 1024 * 100

	tinyCount   = 16384 // 1 MB
	smallCount  = 2048  // 1 MB
	mediumCount = 1024  // 1 MB
	largeCount  = 102   // ~1 MB
	hugeCount   = 10    // ~1 MB
)

type byteSlab struct {
	buffer     []byte
	bufferSize uintptr
	stride     uintptr
	basePtr    *uintptr
	nextPtr    *uintptr
}

// BytePool is a fragmentation friendly way to allocated byte slices.
type BytePool struct {
	tinySlab   byteSlab
	smallSlab  byteSlab
	mediumSlab byteSlab
	largeSlab  byteSlab
	hugeSlab   byteSlab
}

func newByteSlab(size, count int) byteSlab {
	bufferSize := count * size
	buffer := make([]byte, bufferSize)
	basePtr := (*reflect.SliceHeader)(unsafe.Pointer(&buffer)).Data
	nextPtr := basePtr + uintptr(bufferSize)

	return byteSlab{
		buffer:     buffer,
		bufferSize: uintptr(bufferSize),
		stride:     uintptr(size),
		basePtr:    &basePtr,
		nextPtr:    &nextPtr,
	}
}

func (slab *byteSlab) getSlice(size int) (chunk []byte) {
	chunkHeader := (*reflect.SliceHeader)(unsafe.Pointer(&chunk))
	chunkHeader.Len = size
	chunkHeader.Cap = int(slab.stride)

	for {
		// WARNING: The following two lines are order sensitive
		basePtr := atomic.LoadUintptr(slab.basePtr)
		nextPtr := atomic.AddUintptr(slab.nextPtr, -slab.stride)
		lastPtr := basePtr + slab.bufferSize

		switch {
		case nextPtr < basePtr || nextPtr >= lastPtr:
			// out of range either means alloc while realloc or race between
			// base and next during realloc. In the latter case we lose a chunk.
			runtime.Gosched()

		case nextPtr == basePtr:
			// Last item: realloc
			slab.buffer = make([]byte, slab.bufferSize)
			dataPtr := (*reflect.SliceHeader)(unsafe.Pointer(&slab.buffer)).Data

			// WARNING: The following two lines are order sensitive
			atomic.StoreUintptr(slab.nextPtr, dataPtr+slab.bufferSize)
			atomic.StoreUintptr(slab.basePtr, dataPtr)
			fallthrough

		default:
			chunkHeader.Data = nextPtr
			return
		}
	}
}

// NewBytePool creates a new BytePool with each slab using 1 MB of storage.
// The pool contains 5 slabs of different sizes: 64B, 512B, 1KB, 10KB and 100KB.
// Allocations above 100KB will be allocated directly.
func NewBytePool() BytePool {
	return BytePool{
		tinySlab:   newByteSlab(tiny, tinyCount),
		smallSlab:  newByteSlab(small, smallCount),
		mediumSlab: newByteSlab(medium, mediumCount),
		largeSlab:  newByteSlab(large, largeCount),
		hugeSlab:   newByteSlab(huge, hugeCount),
	}
}

// NewBytePoolWithSize creates a new BytePool with each slab size using n MB of
// storage. See NewBytePool() for slab size details.
func NewBytePoolWithSize(n int) BytePool {
	if n <= 0 {
		n = 1
	}
	return BytePool{
		tinySlab:   newByteSlab(tiny, tinyCount*n),
		smallSlab:  newByteSlab(small, smallCount*n),
		mediumSlab: newByteSlab(medium, mediumCount*n),
		largeSlab:  newByteSlab(large, largeCount*n),
		hugeSlab:   newByteSlab(huge, hugeCount*n),
	}
}

// Get returns a slice allocated to a normalized size.
// Sizes are organized in evenly sized buckets so that fragmentation is kept low.
func (b *BytePool) Get(size int) []byte {
	switch {
	case size == 0:
		return []byte{}

	case size <= tiny:
		return b.tinySlab.getSlice(size)

	case size <= small:
		return b.smallSlab.getSlice(size)

	case size <= medium:
		return b.mediumSlab.getSlice(size)

	case size <= large:
		return b.largeSlab.getSlice(size)

	case size <= huge:
		return b.hugeSlab.getSlice(size)

	default:
		return make([]byte, size)
	}
}
<|endoftext|>


<commit_before>package template

import (
	"fmt"
	"path/filepath"
	"runtime"
	"testing"
)

func TestFileSandbox(t *testing.T) {
	// while most of the function can be tested lexigraphically,
	// we need to be able to walk actual symlinks.
	_, filename, _, _ := runtime.Caller(0)
	sandboxDir := filepath.Join(filepath.Dir(filename), "testdata", "sandbox")
	cases := []struct {
		name     string
		sandbox  string
		path     string
		expected error
	}{
		{
			"absolute_path_no_sandbox",
			"",
			"/path/to/file",
			nil,
		},
		{
			"relative_path_no_sandbox",
			"",
			"./path/to/file",
			nil,
		},
		{
			"absolute_path_with_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/file"),
			nil,
		},
		{
			"relative_path_in_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/../to/file"),
			nil,
		},
		{
			"symlink_path_in_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/ok-symlink"),
			nil,
		},
		{
			"relative_path_escaping_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/../../../funcs_test.go"),
			fmt.Errorf("'%s' is outside of sandbox",
				filepath.Join(sandboxDir, "path/../../../funcs_test.go")),
		},
		{
			"symlink_escaping_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/bad-symlink"),
			fmt.Errorf("'%s' is outside of sandbox",
				filepath.Join(sandboxDir, "path/to/bad-symlink")),
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			err := pathInSandbox(tc.sandbox, tc.path)
			if err != nil && tc.expected != nil {
				if err.Error() != tc.expected.Error() {
					t.Fatalf("expected %v got %v", tc.expected, err)
				}
			} else if err != tc.expected {
				t.Fatalf("expected %v got %v", tc.expected, err)
			}
		})
	}
}
<commit_msg>add tests for byMeta function<commit_after>package template

import (
	"fmt"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"

	dep "github.com/hashicorp/consul-template/dependency"
)

func TestFileSandbox(t *testing.T) {
	// while most of the function can be tested lexigraphically,
	// we need to be able to walk actual symlinks.
	_, filename, _, _ := runtime.Caller(0)
	sandboxDir := filepath.Join(filepath.Dir(filename), "testdata", "sandbox")
	cases := []struct {
		name     string
		sandbox  string
		path     string
		expected error
	}{
		{
			"absolute_path_no_sandbox",
			"",
			"/path/to/file",
			nil,
		},
		{
			"relative_path_no_sandbox",
			"",
			"./path/to/file",
			nil,
		},
		{
			"absolute_path_with_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/file"),
			nil,
		},
		{
			"relative_path_in_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/../to/file"),
			nil,
		},
		{
			"symlink_path_in_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/ok-symlink"),
			nil,
		},
		{
			"relative_path_escaping_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/../../../funcs_test.go"),
			fmt.Errorf("'%s' is outside of sandbox",
				filepath.Join(sandboxDir, "path/../../../funcs_test.go")),
		},
		{
			"symlink_escaping_sandbox",
			sandboxDir,
			filepath.Join(sandboxDir, "path/to/bad-symlink"),
			fmt.Errorf("'%s' is outside of sandbox",
				filepath.Join(sandboxDir, "path/to/bad-symlink")),
		},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
			err := pathInSandbox(tc.sandbox, tc.path)
			if err != nil && tc.expected != nil {
				if err.Error() != tc.expected.Error() {
					t.Fatalf("expected %v got %v", tc.expected, err)
				}
			} else if err != tc.expected {
				t.Fatalf("expected %v got %v", tc.expected, err)
			}
		})
	}
}

func Test_byMeta(t *testing.T) {
	svcA := &dep.HealthService{
		Node:                "",
		NodeID:              "",
		NodeAddress:         "",
		NodeTaggedAddresses: nil,
		NodeMeta:            nil,
		ServiceMeta: map[string]string{
			"version":         "v2",
			"version_num":     "2",
			"bad_version_num": "1zz",
			"env":             "dev",
		},
		Address: "",
		ID:      "svcA",
		Name:    "",
		Tags:    nil,
		Checks:  nil,
		Status:  "",
		Port:    0,
	}

	svcB := &dep.HealthService{
		Node:                "",
		NodeID:              "",
		NodeAddress:         "",
		NodeTaggedAddresses: nil,
		NodeMeta:            nil,
		ServiceMeta: map[string]string{
			"version":         "v11",
			"version_num":     "11",
			"bad_version_num": "1zz",
			"env":             "prod",
		},
		Address: "",
		ID:      "svcB",
		Name:    "",
		Tags:    nil,
		Checks:  nil,
		Status:  "",
		Port:    0,
	}

	svcC := &dep.HealthService{
		Node:                "",
		NodeID:              "",
		NodeAddress:         "",
		NodeTaggedAddresses: nil,
		NodeMeta:            nil,
		ServiceMeta: map[string]string{
			"version":         "v11",
			"version_num":     "11",
			"bad_version_num": "1zz",
			"env":             "prod",
		},
		Address: "",
		ID:      "svcC",
		Name:    "",
		Tags:    nil,
		Checks:  nil,
		Status:  "",
		Port:    0,
	}

	type args struct {
		meta     string
		services []*dep.HealthService
	}

	tests := []struct {
		name       string
		args       args
		wantGroups map[string][]*dep.HealthService
		wantErr    bool
	}{
		{
			name: "version string",
			args: args{
				meta:     "version",
				services: []*dep.HealthService{svcA, svcB, svcC},
			},
			wantGroups: map[string][]*dep.HealthService{
				"v11": {svcB, svcC},
				"v2":  {svcA},
			},
			wantErr: false,
		},
		{
			name: "version number",
			args: args{
				meta:     "version_num|int",
				services: []*dep.HealthService{svcA, svcB, svcC},
			},
			wantGroups: map[string][]*dep.HealthService{
				"00011": {svcB, svcC},
				"00002": {svcA},
			},
			wantErr: false,
		},
		{
			name: "bad version number",
			args: args{
				meta:     "bad_version_num|int",
				services: []*dep.HealthService{svcA, svcB, svcC},
			},
			wantGroups: nil,
			wantErr:    true,
		},
		{
			name: "multiple meta",
			args: args{
				meta:     "env,version_num|int,version",
				services: []*dep.HealthService{svcA, svcB, svcC},
			},
			wantGroups: map[string][]*dep.HealthService{
				"dev_00002_v2":   {svcA},
				"prod_00011_v11": {svcB, svcC},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotGroups, err := byMeta(tt.args.meta, tt.args.services)
			if (err != nil) != tt.wantErr {
				t.Errorf("byMeta() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			onlyIDs := func(groups map[string][]*dep.HealthService) (ids map[string]map[string]int) {
				ids = make(map[string]map[string]int)
				for group, svcs := range groups {
					ids[group] = make(map[string]int)
					for _, svc := range svcs {
						ids[group][svc.ID] = 1
					}
				}
				return
			}

			gotIDs := onlyIDs(gotGroups)
			wantIDs := onlyIDs(tt.wantGroups)
			if !reflect.DeepEqual(gotGroups, tt.wantGroups) {
				t.Errorf("byMeta() = %v, want %v", gotIDs, wantIDs)
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"os"
)

type BitmapHeader struct {
	HeaderField uint16
	Size uint32
	_ uint32
	DataAddress uint32
	DIBSize uint32
	Width uint16
	Height uint16
	ColPlanes uint16
	Bpp uint16
}

type Header struct {
	License     [100]byte
	Name        [12]byte
	Version     [8]byte
	Timestamp   [42]byte
	FileSize    uint32
	DirectoryCount uint16
	FileCount   uint16
	Val1       [2]byte //unidentified
	Val2       [2]byte //unidentified
	Val3       [2]byte //unidentified
	Val4       [2]byte //unidentified
	Val5       [2]byte //unidentified
}

type DirectoryInfo struct {
	Name  [4]byte
	Count uint16
	Pos   uint16
}

type FileInfo struct {
	Name [12]byte
	ID uint16 // 0 = Default, 200 = BMP, 1000 = TXT
	Size uint32
	Addr uint32
	Val1 uint8 //unidentified
	Val2 uint8 //unidentified
	Val3 uint8 //unidentified
	Val4 uint8 //unidentified
}

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func insertByte(slice []byte, index uint16, value byte) []byte {
	s_a := slice[:index]
	s_b := slice[index:]
	s_a = append(s_a, value, value)
	s_a = append(s_a, s_b...)
	return s_a
}

func readNumBytes(file *os.File, number int) []byte {
	bytes := make([]byte, number)
	num, err := file.Read(bytes)
	if num != number {
		fmt.Printf("Ran out of bytes! (wanted: %d, got: %d)\n", number, num)
	}
	check(err)
	return bytes
}

func getBuffer(f *os.File, n int) *bytes.Buffer {
	data := readNumBytes(f, n)
	buffer := bytes.NewBuffer(data)
	return buffer
}

func unpackHeader(f *os.File, hdrSize int) *Header {
	hdr := Header{}
	err := binary.Read(getBuffer(f, hdrSize), binary.LittleEndian, &hdr)
	check(err)
	return &hdr
}

func unpackDirectoryList(f *os.File, cnt int) []*DirectoryInfo {
	dir_list := make([]*DirectoryInfo, cnt)
	for i := 0; i < cnt; i++ {
		dir := DirectoryInfo{}
		err := binary.Read(getBuffer(f, 8), binary.LittleEndian, &dir)
		check(err)
		dir_list[i] = &dir
	}
	return dir_list
}

func unpackFileList(f *os.File, cnt int) []*FileInfo {
	file_list := make([]*FileInfo, cnt)
	for i := 0; i < cnt; i++ {
		file := FileInfo{}
		err := binary.Read(getBuffer(f, 26), binary.LittleEndian, &file)
		check(err)
		file_list[i] = &file
	}
	return file_list
}

func unpackFile(f *os.File, file *FileInfo) []byte {
	addr := int64(file.Addr)
	fsize := int(file.Size)
	f.Seek(addr, 0)
	file_data := readNumBytes(f, fsize)

	return file_data
}

func getPalette(f *os.File, dir_list []*DirectoryInfo, files []*FileInfo, s string) []byte {
	for _, dir := range dir_list {
		if string(dir.Name[:3]) == "PAL" {
			fmt.Printf("PAL directory found\n")
			for _, file := range files[dir.Pos:dir.Pos + dir.Count] {
				file_name := string(bytes.Trim(file.Name[:12], "x\000"))
				if file_name == s {
					fmt.Printf("Unpacking palette: %s\n", file_name)
					pal := unpackFile(f, file)

					//PAL values are 0x00 - 0x3F, and red/blue channels seem to be swapped
					for i := 0; i < len(pal); i++ {
						pal[i] *= 4
						if (i+1) % 3 == 0 {
							buf := pal[i]
							pal[i] = pal[i-2]
							pal[i-2] = buf
						}
					}
					return pal
				}
			}
		}
	}
	log.Fatal("Couldn't find requested PAL file")
	return nil
}

func unpackFiles(f *os.File, hdr *Header, dir_list []*DirectoryInfo, files []*FileInfo, pal []byte)  {
	var buf bytes.Buffer

	for _, dir := range dir_list {
		work_dir := "./" +  string(bytes.Trim(hdr.Name[:8], "x\000")) + "/" +
			string(dir.Name[:3]) + "/"
		fmt.Printf("Extracting to %s\n", work_dir)
		os.MkdirAll(work_dir, os.ModePerm)
		fmt.Printf("File count: %d\n", dir.Count)

		for _, file := range files[dir.Pos:dir.Count + dir.Pos] {
			s := work_dir + string(bytes.Trim(file.Name[:12], "x\000"))
			out, err := os.Create(s)
			check(err)

			out_data := unpackFile(f, file)

			//fmt.Printf("Filename: %s\n ID: %x\n Val1: %x\n Val2: %x\n Val3: %x\n Val4: %x\n",
			//	file.Name, file.ID, file.Val1, file.Val2, file.Val3, file.Val4)
			switch file.ID {
			case 0x200: //Bitmap
				dim := out_data[:4]
				bmp_x := binary.LittleEndian.Uint16(dim[:2])
				bmp_y := binary.LittleEndian.Uint16(dim[2:])
				bmp_data := out_data[6:]
				bmp_header := BitmapHeader{
					HeaderField: 0x4d42,
					Size: uint32(0x31D + file.Size),
					DataAddress: 0x31D,
					DIBSize: 0x0c,
					Width: bmp_x,
					Height: bmp_y,
					ColPlanes: 0x1,
					Bpp: 0x8,
				}

				//Some bitmaps are not 4-byte aligned, so we need to add the padding manually
				if bmp_x % 4 != 0 {
					fmt.Printf("File %s requires padding\n", file.Name)
					row := uint16(bmp_x)
					for i := row; i < uint16(len(bmp_data)) - row; i += row {
						bmp_data = insertByte(bmp_data, i, 0)
					}
				}

				binary.Write(&buf, binary.LittleEndian, bmp_header)
				buf.Write(pal)
				binary.Write(&buf, binary.LittleEndian, bmp_data)

				bmp_file := make([]byte, buf.Len())
				err = binary.Read(&buf, binary.LittleEndian, bmp_file)
				check(err)
				_, err = out.Write(bmp_file)
				check(err)

			default:
				_, err = out.Write(out_data)
				check(err)
			}
			out.Close()
		}
	}
}

func main() {
	var hdrSize int

	if len(os.Args) == 1 {
		fmt.Printf("Usage: rsfunpack FILE\n")
		return
	}
	path := os.Args[1]

	f, err := os.Open(path)
	check(err)
	defer f.Close()

	fmt.Printf("%s opened\n", path)

	formatCheck := readNumBytes(f, 1)

	if formatCheck[0] == byte(0x41) {
		fmt.Printf("Valid RSF format found\n")
		hdrSize = 0xb4
	} else if formatCheck[0] == byte(0x6c) {
		log.Fatal("Cannot handle old-style RSF format\n")
	} else {
		log.Fatal("Unknown file format\n")
	}

	f.Seek(0, 0)
	header := unpackHeader(f, hdrSize)

	fmt.Printf("%s\n%s\n%s\n%s\nFilesize: %d\nFormats: %d Files: %d\n", header.License, header.Name,
		header.Version, header.Timestamp, header.FileSize, header.DirectoryCount, header.FileCount)

	directory_list := unpackDirectoryList(f, int(header.DirectoryCount))
	file_list := unpackFileList(f, int(header.FileCount))

	rgb_pal := getPalette(f, directory_list, file_list, "TRUERGB.PAL")
	//l23_pal := getPalette(f, header, format_list, file_list, "L23.PAL")

	unpackFiles(f, header, directory_list, file_list, rgb_pal)
}
<commit_msg>Changed BMP header format<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"os"
)

type RGB struct {
	r uint8
	g uint8
	b uint8
}

type RGBA struct {
	r uint8
	g uint8
	b uint8
	_ uint8
}

type BitmapHeader struct {
	HeaderField uint16
	Size uint32
	_ uint32
	DataAddress uint32
	DIBSize uint32
	Width uint32
	Height uint32
	ColPlanes uint16
	Bpp uint16
	_ [24]byte
}

type Header struct {
	License     [100]byte
	Name        [12]byte
	Version     [8]byte
	Timestamp   [42]byte
	FileSize    uint32
	DirectoryCount uint16
	FileCount   uint16
	Val1       [2]byte //unidentified
	Val2       [2]byte //unidentified
	Val3       [2]byte //unidentified
	Val4       [2]byte //unidentified
	Val5       [2]byte //unidentified
}

type DirectoryInfo struct {
	Name  [4]byte
	Count uint16
	Pos   uint16
}

type FileInfo struct {
	Name [12]byte
	ID uint16 // 0 = Default, 200 = BMP, 1000 = TXT
	Size uint32
	Addr uint32
	Val1 uint8 //unidentified
	Val2 uint8 //unidentified
	Val3 uint8 //unidentified
	Val4 uint8 //unidentified
}

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func insertByte(slice []byte, index int, value byte) []byte {
	s_a := slice[:index+1]
	s_b := slice[index:]
	s_a = append(s_a, value)
	s_a = append(s_a, s_b...)
	return s_a
}

func readNumBytes(file *os.File, number int) []byte {
	bytes := make([]byte, number)
	num, err := file.Read(bytes)
	if num != number {
		fmt.Printf("Ran out of bytes! (wanted: %d, got: %d)\n", number, num)
	}
	check(err)
	return bytes
}

func getBuffer(f *os.File, n int) *bytes.Buffer {
	data := readNumBytes(f, n)
	buffer := bytes.NewBuffer(data)
	return buffer
}

func unpackHeader(f *os.File, hdrSize int) *Header {
	hdr := Header{}
	err := binary.Read(getBuffer(f, hdrSize), binary.LittleEndian, &hdr)
	check(err)
	return &hdr
}

func unpackDirectoryList(f *os.File, cnt int) []*DirectoryInfo {
	dir_list := make([]*DirectoryInfo, cnt)
	for i := 0; i < cnt; i++ {
		dir := DirectoryInfo{}
		err := binary.Read(getBuffer(f, 8), binary.LittleEndian, &dir)
		check(err)
		dir_list[i] = &dir
	}
	return dir_list
}

func unpackFileList(f *os.File, cnt int) []*FileInfo {
	file_list := make([]*FileInfo, cnt)
	for i := 0; i < cnt; i++ {
		file := FileInfo{}
		err := binary.Read(getBuffer(f, 26), binary.LittleEndian, &file)
		check(err)
		file_list[i] = &file
	}
	return file_list
}

func unpackFile(f *os.File, file *FileInfo) []byte {
	addr := int64(file.Addr)
	fsize := int(file.Size)
	f.Seek(addr, 0)
	file_data := readNumBytes(f, fsize)

	return file_data
}

func getPalette(f *os.File, dir_list []*DirectoryInfo, files []*FileInfo, s string) []*RGB {
	for _, dir := range dir_list {
		if string(dir.Name[:3]) == "PAL" {
			fmt.Printf("PAL directory found\n")
			for _, file := range files[dir.Pos:dir.Pos + dir.Count] {
				file_name := string(bytes.Trim(file.Name[:12], "x\000"))
				if file_name == s {
					fmt.Printf("Unpacking palette: %s\n", file_name)
					palette := make([]*RGB, 256)
					f.Seek(int64(file.Addr), 0)
					for i := 0; i < 256; i++ {
						pal := readNumBytes(f, 3)
						pal_entry := RGB{
							r : pal[0],
							g : pal[1],
							b : pal[2],
						}
						fmt.Printf("%+x\n",pal_entry)
						palette[i] = &pal_entry
					}
					return palette
				}
			}
		}
	}
	log.Fatal("Couldn't find requested PAL file")
	return nil
}

func unpackFiles(f *os.File, hdr *Header, dir_list []*DirectoryInfo, files []*FileInfo, pal []*RGB)  {
	var buf bytes.Buffer

	for _, dir := range dir_list {
		work_dir := "./" +  string(bytes.Trim(hdr.Name[:8], "x\000")) + "/" +
			string(dir.Name[:3]) + "/"
		fmt.Printf("Extracting to %s\n", work_dir)
		os.MkdirAll(work_dir, os.ModePerm)
		fmt.Printf("File count: %d\n", dir.Count)

		for _, file := range files[dir.Pos:dir.Count + dir.Pos] {
			s := work_dir + string(bytes.Trim(file.Name[:12], "x\000"))
			out, err := os.Create(s)
			check(err)

			out_data := unpackFile(f, file)

			//fmt.Printf("Filename: %s\n ID: %x\n Val1: %x\n Val2: %x\n Val3: %x\n Val4: %x\n",
			//	file.Name, file.ID, file.Val1, file.Val2, file.Val3, file.Val4)
			switch file.ID {
			case 0x200: //Bitmap
				dim := out_data[:4]
				bmp_x := uint32(binary.LittleEndian.Uint16(dim[:2]))
				bmp_y := uint32(binary.LittleEndian.Uint16(dim[2:]))
				bmp_data := out_data[4:]
				bmp_header := BitmapHeader{
					HeaderField: 0x4d42,
					Size: uint32(0x43B + file.Size),
					DataAddress: 0x43B,
					DIBSize: 0x28,
					Width: bmp_x,
					Height: bmp_y,
					ColPlanes: 0x1,
					Bpp: 0x8,
				}

				//Some bitmaps are not 4-byte aligned, so we need to pad them manually
				row := int(bmp_x)
				if row % 4 != 0 {
					fmt.Printf("File %s requires padding\n", file.Name)
					for i := row; i < len(bmp_data) - row; i += row {
						for ii := 1; ii < row % 4; ii++ {
							bmp_data = insertByte(bmp_data, i, 0)
						}
					}
				}

				binary.Write(&buf, binary.LittleEndian, bmp_header)

				//PAL values are 0x00 - 0x3F, and red/blue channels seem to be swapped
				for i := 0; i < len(pal); i++ {
					outpal_entry := RGBA{
						r : pal[i].b * 4,
						g : pal[i].g * 4,
						b : pal[i].r * 4,
					}
					binary.Write(&buf, binary.LittleEndian, outpal_entry)
				}

				binary.Write(&buf, binary.LittleEndian, bmp_data)

				bmp_file := make([]byte, buf.Len())
				err = binary.Read(&buf, binary.LittleEndian, bmp_file)
				check(err)
				_, err = out.Write(bmp_file)
				check(err)

			default:
				_, err = out.Write(out_data)
				check(err)
			}
			out.Close()
		}
	}
}

func main() {
	var hdrSize int

	if len(os.Args) == 1 {
		fmt.Printf("Usage: rsfunpack FILE\n")
		return
	}
	path := os.Args[1]

	f, err := os.Open(path)
	check(err)
	defer f.Close()

	fmt.Printf("%s opened\n", path)

	formatCheck := readNumBytes(f, 1)

	if formatCheck[0] == byte(0x41) {
		fmt.Printf("Valid RSF format found\n")
		hdrSize = 0xb4
	} else if formatCheck[0] == byte(0x6c) {
		log.Fatal("Cannot handle old-style RSF format\n")
	} else {
		log.Fatal("Unknown file format\n")
	}

	f.Seek(0, 0)
	header := unpackHeader(f, hdrSize)

	fmt.Printf("%s\n%s\n%s\n%s\nFilesize: %d\nFormats: %d Files: %d\n", header.License, header.Name,
		header.Version, header.Timestamp, header.FileSize, header.DirectoryCount, header.FileCount)

	directory_list := unpackDirectoryList(f, int(header.DirectoryCount))
	file_list := unpackFileList(f, int(header.FileCount))

	rgb_pal := getPalette(f, directory_list, file_list, "TRUERGB.PAL")
	//l23_pal := getPalette(f, header, format_list, file_list, "L23.PAL")

	unpackFiles(f, header, directory_list, file_list, rgb_pal)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/adamar/zego"
	"log"
)

func main() {

	auth := zego.Auth{"user@example.com", "password", "subdomain.zendesk.com"}
        response := auth.ListTickets()
        tickets := &zego.TicketArray{}
        json.Unmarshal([]byte(response.Raw), tickets)
	log.Print(tickets)

}
<commit_msg>Bug fixes<commit_after>package main

import (
	"github.com/adamar/zego/zego"
	"log"
	"encoding/json"
)

func main() {
	auth := zego.Auth{"user@example.com", "Password", "example.zendesk.com"}
	// auth := zego.Auth{"user@example.com/token", "Token", "example.zendesk.com"}
	response, _ := auth.ListTickets()
	tickets := &zego.TicketArray{}
	json.Unmarshal([]byte(response.Raw), tickets)
	log.Print(tickets)
}
<|endoftext|>


<commit_before>// This is an example program that shows the usage of the memsearch package.
//
// With this program you can:
//   - Search for a string in the memory of a process with a given PID
//   - Print an arbitrary amount of bytes from the process memory.
package main

import (
	"encoding/hex"
	"flag"
	"io/ioutil"
	"log"
	"regexp"
	"strings"

	"github.com/mozilla/masche/memaccess"
	"github.com/mozilla/masche/memsearch"
	"github.com/mozilla/masche/process"
)

var (
	action = flag.String("action", "<nil>", "Action to perfom. One of: search, regexp-search, file-search, print")
	pid    = flag.Int("pid", 0, "Process id to analyze")
	addr   = flag.Int("addr", 0x0, "The initial address in the process address space to search/print")

	// print action flags
	size = flag.Int("n", 4, "Amount of bytes to print")

	// search action flags
	needle = flag.String("needle", "Find This!", "String to search for (interpreted as []byte)")

	// regexp-search action flags
	regexpString = flag.String("regexp", "regexp?", "Regexp to search for")

	// file-search action flags
	fileneedle = flag.String("fileneedle", "example.in", "Filename that contains hex-encoded needle (spaces are ignored)")
)

func logErrors(harderror error, softerrors []error) {
	if harderror != nil {
		log.Fatal(harderror)
	}
	for _, soft := range softerrors {
		log.Print(soft)
	}
}

func main() {
	flag.Parse()

	proc, harderror, softerrors := process.OpenFromPid(uint(*pid))
	logErrors(harderror, softerrors)

	switch *action {

	case "<nil>":
		log.Fatal("Missing action flag.")
	case "file-search":
		data, err := ioutil.ReadFile(*fileneedle)
		if err != nil {
			log.Fatal(err)
		}
		encoded := strings.Replace(strings.TrimSpace(string(data)), " ", "", -1)
		data, err = hex.DecodeString(encoded)
		if err != nil {
			log.Fatal(err)
		}
		found, address, harderror, softerrors := memsearch.FindBytesSequence(proc, uintptr(*addr), data)
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "search":
		found, address, harderror, softerrors := memsearch.FindBytesSequence(proc, uintptr(*addr), []byte(*needle))
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "regexp-search":
		r, err := regexp.Compile(*regexpString)
		if err != nil {
			log.Fatal(err)
		}

		found, address, harderror, softerrors := memsearch.FindRegexpMatch(proc, uintptr(*addr), r)
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "print":
		buf := make([]byte, *size)
		harderror, softerrors = memaccess.CopyMemory(proc, uintptr(*addr), buf)
		logErrors(harderror, softerrors)
		log.Println(string(buf))

	default:
		log.Fatal("Unrecognized action ", *action)
	}
}
<commit_msg>[examples] Ignore newlines in memsearch file-search<commit_after>// This is an example program that shows the usage of the memsearch package.
//
// With this program you can:
//   - Search for a string in the memory of a process with a given PID
//   - Print an arbitrary amount of bytes from the process memory.
package main

import (
	"encoding/hex"
	"flag"
	"io/ioutil"
	"log"
	"regexp"
	"strings"

	"github.com/mozilla/masche/memaccess"
	"github.com/mozilla/masche/memsearch"
	"github.com/mozilla/masche/process"
)

var (
	action = flag.String("action", "<nil>", "Action to perfom. One of: search, regexp-search, file-search, print")
	pid    = flag.Int("pid", 0, "Process id to analyze")
	addr   = flag.Int("addr", 0x0, "The initial address in the process address space to search/print")

	// print action flags
	size = flag.Int("n", 4, "Amount of bytes to print")

	// search action flags
	needle = flag.String("needle", "Find This!", "String to search for (interpreted as []byte)")

	// regexp-search action flags
	regexpString = flag.String("regexp", "regexp?", "Regexp to search for")

	// file-search action flags
	fileneedle = flag.String("fileneedle", "example.in", "Filename that contains hex-encoded needle (spaces are ignored)")
)

func logErrors(harderror error, softerrors []error) {
	if harderror != nil {
		log.Fatal(harderror)
	}
	for _, soft := range softerrors {
		log.Print(soft)
	}
}

func main() {
	flag.Parse()

	proc, harderror, softerrors := process.OpenFromPid(uint(*pid))
	logErrors(harderror, softerrors)

	switch *action {

	case "<nil>":
		log.Fatal("Missing action flag.")
	case "file-search":
		data, err := ioutil.ReadFile(*fileneedle)
		if err != nil {
			log.Fatal(err)
		}
		encoded := strings.Replace(strings.Replace(strings.TrimSpace(string(data)), " ", "", -1), "\n", "", -1)
		data, err = hex.DecodeString(encoded)
		if err != nil {
			log.Fatal(err)
		}
		found, address, harderror, softerrors := memsearch.FindBytesSequence(proc, uintptr(*addr), data)
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "search":
		found, address, harderror, softerrors := memsearch.FindBytesSequence(proc, uintptr(*addr), []byte(*needle))
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "regexp-search":
		r, err := regexp.Compile(*regexpString)
		if err != nil {
			log.Fatal(err)
		}

		found, address, harderror, softerrors := memsearch.FindRegexpMatch(proc, uintptr(*addr), r)
		logErrors(harderror, softerrors)
		if found {
			log.Printf("Found in address: %x\n", address)
		}

	case "print":
		buf := make([]byte, *size)
		harderror, softerrors = memaccess.CopyMemory(proc, uintptr(*addr), buf)
		logErrors(harderror, softerrors)
		log.Println(string(buf))

	default:
		log.Fatal("Unrecognized action ", *action)
	}
}
<|endoftext|>


<commit_before>package events_test

import (
	"sync"
	"sync/atomic"

	"github.com/30x/apid-core"
	"github.com/30x/apid-core/events"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"strconv"
)

var _ = Describe("Events Service", func() {

	It("should ignore event with no listeners", func() {
		em := events.CreateService()
		defer em.Close()
		em.Emit("no listeners", &test_event{"test"})
	})

	It("should publish an event to a listener", func(done Done) {
		em := events.CreateService()

		h := test_handler{
			"handler",
			func(event apid.Event) {
				defer GinkgoRecover()

				em.Close()
				close(done)
			},
		}

		em.Listen("selector", &h)
		em.Emit("selector", &test_event{"test"})
	})

	It("should publish an event to a listener func", func(done Done) {
		em := events.CreateService()

		h := func(event apid.Event) {
			defer GinkgoRecover()

			em.Close()
			close(done)
		}

		em.ListenFunc("selector", h)
		em.Emit("selector", &test_event{"test"})
	})

	It("should publish multiple events to a listener", func(done Done) {
		em := events.CreateService()

		count := int32(0)
		h := test_handler{
			"handler",
			func(event apid.Event) {
				defer GinkgoRecover()

				c := atomic.AddInt32(&count, 1)
				if c > 1 {
					em.Close()
					close(done)
				}
			},
		}

		em.Listen("selector", &h)
		em.Emit("selector", &test_event{"test1"})
		em.Emit("selector", &test_event{"test2"})
	})

	It("EmitWithCallback should call the callback when done with delivery", func(done Done) {
		em := events.CreateService()

		delivered := func(event apid.Event) {
			defer GinkgoRecover()
			close(done)
		}

		em.EmitWithCallback("selector", &test_event{"test1"}, delivered)
	})

	It("should publish only one event to a listenOnce", func(done Done) {
		em := events.CreateService()

		count := 0
		h := func(event apid.Event) {
			defer GinkgoRecover()
			count++
		}

		delivered := func(event apid.Event) {
			defer GinkgoRecover()
			Expect(count).To(Equal(1))
			em.Close()
			close(done)
		}

		em.ListenOnceFunc("selector", h)
		em.Emit("selector", &test_event{"test1"})
		em.EmitWithCallback("selector", &test_event{"test2"}, delivered)
	})

	It("should publish an event to multiple listeners", func(done Done) {
		em := events.CreateService()
		defer em.Close()

		mut := sync.Mutex{}
		hitH1 := false
		hitH2 := false
		h1 := test_handler{
			"handler 1",
			func(event apid.Event) {
				defer GinkgoRecover()
				mut.Lock()
				defer mut.Unlock()
				hitH1 = true
				if hitH1 && hitH2 {
					em.Close()
					close(done)
				}
			},
		}
		h2 := test_handler{
			"handler 2",
			func(event apid.Event) {
				defer GinkgoRecover()

				mut.Lock()
				defer mut.Unlock()
				hitH2 = true
				if hitH1 && hitH2 {
					em.Close()
					close(done)
				}
			},
		}

		em.Listen("selector", &h1)
		em.Listen("selector", &h2)
		em.Emit("selector", &test_event{"test"})
	})

	It("should publish an event delivered event", func(done Done) {
		em := events.CreateService()
		testEvent := &test_event{"test"}
		var testSelector apid.EventSelector = "selector"

		dummy := func(event apid.Event) {}
		em.ListenFunc(testSelector, dummy)

		h := test_handler{
			"event delivered handler",
			func(event apid.Event) {
				defer GinkgoRecover()

				e, ok := event.(apid.EventDeliveryEvent)

				Expect(ok).To(BeTrue())
				Expect(e.Event).To(Equal(testEvent))
				Expect(e.Selector).To(Equal(testSelector))

				em.Close()
				close(done)
			},
		}

		em.Listen(apid.EventDeliveredSelector, &h)
		em.Emit(testSelector, testEvent)
	})

	It("should be able to remove a listener", func(done Done) {
		em := events.CreateService()

		event1 := &test_event{"test1"}
		event2 := &test_event{"test2"}
		event3 := &test_event{"test3"}

		dummy := func(event apid.Event) {}
		em.ListenFunc("selector", dummy)

		h := test_handler{
			"handler",
			func(event apid.Event) {
				defer GinkgoRecover()

				Expect(event).NotTo(Equal(event2))
				if event == event3 {
					em.Close()
					close(done)
				}
			},
		}
		em.Listen("selector", &h)

		// need to drive test like this because of async delivery
		td := test_handler{
			"test driver",
			func(event apid.Event) {
				defer GinkgoRecover()

				e := event.(apid.EventDeliveryEvent)
				if e.Event == event1 {
					em.StopListening("selector", &h)
					em.Emit("selector", event2)
				} else if e.Event == event2 {
					em.Listen("selector", &h)
					em.Emit("selector", event3)
				}
			},
		}
		em.Listen(apid.EventDeliveredSelector, &td)

		em.Emit("selector", event1)
	})

	It("should deliver events according selector", func(done Done) {
		em := events.CreateService()

		e1 := &test_event{"test1"}
		e2 := &test_event{"test2"}

		count := int32(0)

		h1 := test_handler{
			"handler1",
			func(event apid.Event) {
				defer GinkgoRecover()

				c := atomic.AddInt32(&count, 1)
				Expect(event).Should(Equal(e1))
				if c == 2 {
					em.Close()
					close(done)
				}
			},
		}

		h2 := test_handler{
			"handler2",
			func(event apid.Event) {
				defer GinkgoRecover()

				c := atomic.AddInt32(&count, 1)
				Expect(event).Should(Equal(e2))
				if c == 2 {
					em.Close()
					close(done)
				}
			},
		}

		em.Listen("selector1", &h1)
		em.Listen("selector2", &h2)

		em.Emit("selector2", e2)
		em.Emit("selector1", e1)
	})

	It("should publish PluginsInitialized event", func(done Done) {
		xData := make(map[string]interface{})
		xData["schemaVersion"] = "1.2.3"
		p := func(s apid.Services) (pd apid.PluginData, err error) {
			pd = apid.PluginData{
				Name:      "test plugin",
				Version:   "1.0.0",
				ExtraData: xData,
			}
			return
		}
		apid.RegisterPlugin(p)

		h := func(event apid.Event) {
			defer GinkgoRecover()

			if pie, ok := event.(apid.PluginsInitializedEvent); ok {

				apid.Events().Close()

				Expect(len(pie.Plugins)).Should(Equal(1))
				p := pie.Plugins[0]
				Expect(p.Name).To(Equal("test plugin"))
				Expect(p.Version).To(Equal("1.0.0"))
				Expect(p.ExtraData["schemaVersion"]).To(Equal("1.2.3"))

				close(done)
			}
		}
		apid.Events().ListenFunc(apid.SystemEventsSelector, h)

		apid.InitializePlugins("")
	})

	It("shutdown event should be emitted and listened successfully", func(done Done) {
		h := func(event apid.Event) {
			defer GinkgoRecover()

			if pie, ok := event.(apid.ShutdownEvent); ok {
				apid.Events().Close()
				Expect(pie.Description).Should(Equal("apid is going to shutdown"))
				fmt.Println(pie.Description)
				close(done)
			} else {
				Fail("Received wrong event")
			}
		}
		apid.Events().ListenFunc(apid.ShutdownEventSelector, h)
		shutdownHandler := func(event apid.Event) {}
		apid.Events().EmitWithCallback(apid.ShutdownEventSelector, apid.ShutdownEvent{"apid is going to shutdown"}, shutdownHandler)
	})

	It("handlers registered by plugins should execute before apid shutdown", func(done Done) {
		pluginNum := 10
		count := int32(0)

		// create and register plugins, listen to shutdown event
		for i:=0; i<pluginNum; i++ {
			apid.RegisterPlugin(createDummyPlugin(i))
			h := func(event apid.Event) {
				if pie, ok := event.(apid.ShutdownEvent); ok {
					Expect(pie.Description).Should(Equal("apid is going to shutdown"))
					atomic.AddInt32(&count, 1)
				} else {
					Fail("Received wrong event")
				}
			}
			apid.Events().ListenFunc(apid.ShutdownEventSelector, h)
		}

		apid.InitializePlugins("")

		apid.ShutdownPluginsAndWait()

		defer GinkgoRecover()
		apid.Events().Close()

		// handlers of all registered plugins have executed
		Expect(count).Should(Equal(int32(pluginNum)))

		close(done)
	})

	It("should be able to read apid version from PluginsInitialized event", func(done Done) {
		xData := make(map[string]interface{})
		xData["schemaVersion"] = "1.2.3"
		p := func(s apid.Services) (pd apid.PluginData, err error) {
			pd = apid.PluginData{
				Name:      "test plugin",
				Version:   "1.0.0",
				ExtraData: xData,
			}
			return
		}
		apid.RegisterPlugin(p)

		apidVersion := "dummy_version"

		h := func(event apid.Event) {
			defer GinkgoRecover()

			if pie, ok := event.(apid.PluginsInitializedEvent); ok {

				apid.Events().Close()
				Expect(pie.ApidVersion).To(Equal(apidVersion))
				close(done)
			}
		}
		apid.Events().ListenFunc(apid.SystemEventsSelector, h)

		apid.InitializePlugins(apidVersion)
	})
})

func createDummyPlugin(id int) apid.PluginInitFunc{
	xData := make(map[string]interface{})
	xData["schemaVersion"] = "1.2.3"
	p := func(s apid.Services) (pd apid.PluginData, err error) {
		pd = apid.PluginData{
			Name: "test plugin " + strconv.Itoa(id),
			Version: "1.0.0",
			ExtraData: xData,
		}
		return
	}
	return p

}

type test_handler struct {
	description string
	f           func(event apid.Event)
}

func (t *test_handler) String() string {
	return t.description
}

func (t *test_handler) Handle(event apid.Event) {
	t.f(event)
}

type test_event struct {
	description string
}
<commit_msg>fix test case<commit_after>package events_test

import (
	"sync"
	"sync/atomic"

	"github.com/30x/apid-core"
	"github.com/30x/apid-core/events"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"strconv"
)

var _ = Describe("Events Service", func() {

	Context("local", func() {

		var em apid.EventsService

		BeforeEach(func() {
			em = events.CreateService()
		})

		AfterEach(func() {
			if em != nil {
				em.Close()
				em = nil
			}
		})

		It("should ignore event with no listeners", func() {
			em.Emit("no listeners", &test_event{"test"})
		})

		It("should publish an event to a listener", func(done Done) {
			h := test_handler{
				"handler",
				func(event apid.Event) {
					close(done)
				},
			}

			em.Listen("selector", &h)
			em.Emit("selector", &test_event{"test"})
		})

		It("should publish an event to a listener func", func(done Done) {
			h := func(event apid.Event) {
				close(done)
			}

			em.ListenFunc("selector", h)
			em.Emit("selector", &test_event{"test"})
		})

		It("should publish multiple events to a listener", func(done Done) {
			count := int32(0)
			h := test_handler{
				"handler",
				func(event apid.Event) {
					defer GinkgoRecover()

					c := atomic.AddInt32(&count, 1)
					if c > 1 {
						close(done)
					}
				},
			}

			em.Listen("selector", &h)
			em.Emit("selector", &test_event{"test1"})
			em.Emit("selector", &test_event{"test2"})
		})

		It("EmitWithCallback should call the callback when done with delivery", func(done Done) {
			delivered := func(event apid.Event) {
				close(done)
			}

			em.EmitWithCallback("selector", &test_event{"test1"}, delivered)
		})

		It("should publish only one event to a listenOnce", func(done Done) {
			count := 0
			h := func(event apid.Event) {
				defer GinkgoRecover()
				count++
			}

			delivered := func(event apid.Event) {
				defer GinkgoRecover()
				Expect(count).To(Equal(1))
				close(done)
			}

			em.ListenOnceFunc("selector", h)
			em.Emit("selector", &test_event{"test1"})
			em.EmitWithCallback("selector", &test_event{"test2"}, delivered)
		})

		It("should publish an event to multiple listeners", func(done Done) {
			mut := sync.Mutex{}
			hitH1 := false
			hitH2 := false
			h1 := test_handler{
				"handler 1",
				func(event apid.Event) {
					defer GinkgoRecover()
					mut.Lock()
					defer mut.Unlock()
					hitH1 = true
					if hitH1 && hitH2 {
						close(done)
					}
				},
			}
			h2 := test_handler{
				"handler 2",
				func(event apid.Event) {
					defer GinkgoRecover()
					mut.Lock()
					defer mut.Unlock()
					hitH2 = true
					if hitH1 && hitH2 {
						close(done)
					}
				},
			}

			em.Listen("selector", &h1)
			em.Listen("selector", &h2)
			em.Emit("selector", &test_event{"test"})
		})

		It("should publish an event delivered event", func(done Done) {
			testEvent := &test_event{"test"}
			var testSelector apid.EventSelector = "selector"

			dummy := func(event apid.Event) {}
			em.ListenFunc(testSelector, dummy)

			h := test_handler{
				"event delivered handler",
				func(event apid.Event) {
					defer GinkgoRecover()

					e, ok := event.(apid.EventDeliveryEvent)

					Expect(ok).To(BeTrue())
					Expect(e.Event).To(Equal(testEvent))
					Expect(e.Selector).To(Equal(testSelector))

					close(done)
				},
			}

			em.Listen(apid.EventDeliveredSelector, &h)
			em.Emit(testSelector, testEvent)
		})

		It("should be able to remove a listener", func(done Done) {
			event1 := &test_event{"test1"}
			event2 := &test_event{"test2"}
			event3 := &test_event{"test3"}

			dummy := func(event apid.Event) {}
			em.ListenFunc("selector", dummy)

			h := test_handler{
				"handler",
				func(event apid.Event) {
					defer GinkgoRecover()

					Expect(event).NotTo(Equal(event2))
					if event == event3 {
						close(done)
					}
				},
			}
			em.Listen("selector", &h)

			// need to drive test like this because of async delivery
			var td apid.EventHandler
			td = &test_handler{
				"test driver",
				func(event apid.Event) {
					defer GinkgoRecover()

					e := event.(apid.EventDeliveryEvent)
					if e.Event == event1 {
						em.StopListening("selector", &h)
						em.Emit("selector", event2)
					} else if e.Event == event2 {
						em.StopListening(apid.EventDeliveredSelector, td)
						em.Listen("selector", &h)
						em.Emit("selector", event3)
					}
				},
			}
			em.Listen(apid.EventDeliveredSelector, td)

			em.Emit("selector", event1)
		})

		It("should deliver events according selector", func(done Done) {
			e1 := &test_event{"test1"}
			e2 := &test_event{"test2"}

			count := int32(0)

			h1 := test_handler{
				"handler1",
				func(event apid.Event) {
					defer GinkgoRecover()

					c := atomic.AddInt32(&count, 1)
					Expect(event).Should(Equal(e1))
					if c == 2 {
						close(done)
					}
				},
			}

			h2 := test_handler{
				"handler2",
				func(event apid.Event) {
					defer GinkgoRecover()

					c := atomic.AddInt32(&count, 1)
					Expect(event).Should(Equal(e2))
					if c == 2 {
						close(done)
					}
				},
			}

			em.Listen("selector1", &h1)
			em.Listen("selector2", &h2)

			em.Emit("selector2", e2)
			em.Emit("selector1", e1)
		})
	})

	Context("plugins", func() {

		BeforeEach(func() {
		})

		AfterEach(func() {
			apid.Events().Close()
		})

		It("should publish PluginsInitialized event", func(done Done) {
			xData := make(map[string]interface{})
			xData["schemaVersion"] = "1.2.3"
			p := func(s apid.Services) (pd apid.PluginData, err error) {
				pd = apid.PluginData{
					Name:      "test plugin",
					Version:   "1.0.0",
					ExtraData: xData,
				}
				return
			}
			apid.RegisterPlugin(p)

			h := func(event apid.Event) {
				defer GinkgoRecover()

				if pie, ok := event.(apid.PluginsInitializedEvent); ok {

					Expect(len(pie.Plugins)).Should(Equal(1))
					p := pie.Plugins[0]
					Expect(p.Name).To(Equal("test plugin"))
					Expect(p.Version).To(Equal("1.0.0"))
					Expect(p.ExtraData["schemaVersion"]).To(Equal("1.2.3"))

					close(done)
				}
			}
			apid.Events().ListenFunc(apid.SystemEventsSelector, h)

			apid.InitializePlugins("")
		})

		It("shutdown event should be emitted and listened successfully", func(done Done) {
			h := func(event apid.Event) {
				defer GinkgoRecover()

				if pie, ok := event.(apid.ShutdownEvent); ok {
					Expect(pie.Description).Should(Equal("apid is going to shutdown"))
					close(done)
				} else {
					Fail("Received wrong event")
				}
			}
			apid.Events().ListenFunc(apid.ShutdownEventSelector, h)
			shutdownHandler := func(event apid.Event) {}
			apid.Events().EmitWithCallback(apid.ShutdownEventSelector, apid.ShutdownEvent{"apid is going to shutdown"}, shutdownHandler)
		})

		It("handlers registered by plugins should execute before apid shutdown", func(done Done) {
			pluginNum := 10
			count := int32(0)

			// create and register plugins, listen to shutdown event
			for i:=0; i<pluginNum; i++ {
				apid.RegisterPlugin(createDummyPlugin(i))
				h := func(event apid.Event) {
					if pie, ok := event.(apid.ShutdownEvent); ok {
						Expect(pie.Description).Should(Equal("apid is going to shutdown"))
						atomic.AddInt32(&count, 1)
					} else {
						Fail("Received wrong event")
					}
				}
				apid.Events().ListenFunc(apid.ShutdownEventSelector, h)
			}

			apid.InitializePlugins("")

			apid.ShutdownPluginsAndWait()

			// handlers of all registered plugins have executed
			Expect(count).Should(Equal(int32(pluginNum)))
			close(done)
		})

		It("should be able to read apid version from PluginsInitialized event", func(done Done) {
			xData := make(map[string]interface{})
			xData["schemaVersion"] = "1.2.3"
			p := func(s apid.Services) (pd apid.PluginData, err error) {
				pd = apid.PluginData{
					Name:      "test plugin",
					Version:   "1.0.0",
					ExtraData: xData,
				}
				return
			}
			apid.RegisterPlugin(p)

			apidVersion := "dummy_version"

			h := func(event apid.Event) {
				defer GinkgoRecover()

				if pie, ok := event.(apid.PluginsInitializedEvent); ok {
					Expect(pie.ApidVersion).To(Equal(apidVersion))
					close(done)
				}
			}
			apid.Events().ListenFunc(apid.SystemEventsSelector, h)

			apid.InitializePlugins(apidVersion)
		})
	})
})

func createDummyPlugin(id int) apid.PluginInitFunc{
	xData := make(map[string]interface{})
	xData["schemaVersion"] = "1.2.3"
	p := func(s apid.Services) (pd apid.PluginData, err error) {
		pd = apid.PluginData{
			Name: "test plugin " + strconv.Itoa(id),
			Version: "1.0.0",
			ExtraData: xData,
		}
		return
	}
	return p
}

type test_handler struct {
	description string
	f           func(event apid.Event)
}

func (t *test_handler) String() string {
	return t.description
}

func (t *test_handler) Handle(event apid.Event) {
	t.f(event)
}

type test_event struct {
	description string
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/amlun/linda"
)

// RPUSH print "{\"queue\":\"print\",\"Payload\":{\"class\":\"printArgs\",\"args\":[\"a\",\"b\",\"c\"]}}"
// RPUSH print "{\"queue\":\"print\",\"Payload\":{\"class\":\"printArgs\",\"args\":[\"A\",\"B\",\"C\"]}}"
// RPUSH print "{\"queue\":\"print\",\"Payload\":{\"class\":\"printArgs\",\"args\":[1,2,3,4,5,6,7]}}"
func init() {
	linda.RegisterWorkers("PrintArgs", PrintArgs)
}

// go run print_args -queue=print -connection=redis://localhost:6379/
func main() {
	if err := linda.Run(); err != nil {
		fmt.Println("Error:", err)
	}
}

func PrintArgs(args ...interface{}) error {
	fmt.Println(args)
	return nil
}
<commit_msg>remove printArgs<commit_after><|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/pkg/browser"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage:\n  %s [file]\n", os.Args[0])
	flag.PrintDefaults()
}

func init() {
	flag.Usage = usage
	flag.Parse()
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	args := flag.Args()
	switch len(args) {
	case 0:
		check(browser.OpenReader(os.Stdin))
	case 1:
		check(browser.OpenFile(args[0]))
	default:
		usage()
	}
}
<commit_msg>added examples<commit_after>// Open is a simple example of the github.com/pkg/browser package.
//
// Usage:
//
//    # Open a file in a browser window
//    Open $FILE
//
//    # Open a URL in a browser window
//    Open $URL
//
//    # Open the contents of stdin in a browser window
//    cat $SOMEFILE | Open
package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/pkg/browser"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage:\n  %s [file]\n", os.Args[0])
	flag.PrintDefaults()
}

func init() {
	flag.Usage = usage
	flag.Parse()
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	args := flag.Args()
	switch len(args) {
	case 0:
		check(browser.OpenReader(os.Stdin))
	case 1:
		check(browser.OpenFile(args[0]))
	default:
		usage()
	}
}
<|endoftext|>


<commit_before>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	// cookieData := ctx.response.Header().Get("Set-Cookie")
	// cookieData += "; SameSite=lax"
	// ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetRequestHeader set the value for the request header.
func (ctx *Context) SetRequestHeader(header string, value string) {
	ctx.request.Header.Set(header, value)
}

// GetResponseHeader sets response header to value.
func (ctx *Context) GetResponseHeader(header string) string {
	return ctx.response.Header().Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RequestBodyJSON returns the JSON parsed request body as map[string]interface{} or []interface{}.
func (ctx *Context) RequestBodyJSON() (interface{}, error) {
	var data interface{}
	err := json.Unmarshal(ctx.RequestBody(), &data)
	return data, err
}

// RequestBodyJSONObject returns the JSON parsed request body as map[string]interface{}.
func (ctx *Context) RequestBodyJSONObject() (map[string]interface{}, error) {
	json, err := ctx.RequestBodyJSON()

	if err != nil {
		return nil, err
	}

	data, formatOK := json.(map[string]interface{})

	if !formatOK {
		return nil, errors.New("Invalid format: Expected JSON object")
	}

	return data, nil
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// IsMediaResponse returns whether the given context has already set its content type to a media type.
func (ctx *Context) IsMediaResponse() bool {
	contentType := ctx.response.Header().Get(contentTypeHeader)
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")
}

// respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) respond(code string) {
	ctx.respondBytes(StringToBytesUnsafe(code))
}

// respondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) respondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	isMedia := ctx.IsMediaResponse()

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<commit_msg>Added context.Data<commit_after>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Custom data
	Data interface{}

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	// cookieData := ctx.response.Header().Get("Set-Cookie")
	// cookieData += "; SameSite=lax"
	// ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetRequestHeader set the value for the request header.
func (ctx *Context) SetRequestHeader(header string, value string) {
	ctx.request.Header.Set(header, value)
}

// GetResponseHeader sets response header to value.
func (ctx *Context) GetResponseHeader(header string) string {
	return ctx.response.Header().Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RequestBodyJSON returns the JSON parsed request body as map[string]interface{} or []interface{}.
func (ctx *Context) RequestBodyJSON() (interface{}, error) {
	var data interface{}
	err := json.Unmarshal(ctx.RequestBody(), &data)
	return data, err
}

// RequestBodyJSONObject returns the JSON parsed request body as map[string]interface{}.
func (ctx *Context) RequestBodyJSONObject() (map[string]interface{}, error) {
	json, err := ctx.RequestBodyJSON()

	if err != nil {
		return nil, err
	}

	data, formatOK := json.(map[string]interface{})

	if !formatOK {
		return nil, errors.New("Invalid format: Expected JSON object")
	}

	return data, nil
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// IsMediaResponse returns whether the given context has already set its content type to a media type.
func (ctx *Context) IsMediaResponse() bool {
	contentType := ctx.response.Header().Get(contentTypeHeader)
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")
}

// respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) respond(code string) {
	ctx.respondBytes(StringToBytesUnsafe(code))
}

// respondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) respondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	isMedia := ctx.IsMediaResponse()

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<|endoftext|>


<commit_before>package vlog_test

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"testing"

	"veyron.io/veyron/veyron/lib/modules"
	"veyron.io/veyron/veyron/lib/testutil"

	"veyron.io/veyron/veyron2/vlog"
)

func TestHelperProcess(t *testing.T) {
	modules.DispatchInTest()
}

func init() {
	testutil.Init()
	modules.RegisterChild("child", "", child)
}

func child(stdin io.Reader, stdout, stderr io.Writer, env map[string]string, args ...string) error {
	tmp := filepath.Join(os.TempDir(), "foo")
	flag.Set("log_dir", tmp)
	flag.Set("vmodule", "foo=2")
	flags := vlog.Log.ExplicitlySetFlags()
	if v, ok := flags["log_dir"]; !ok || v != tmp {
		return fmt.Errorf("log_dir was supposed to be %v", tmp)
	}
	if v, ok := flags["vmodule"]; !ok || v != "foo=2" {
		return fmt.Errorf("vmodule was supposed to be foo=2")
	}
	if f := flag.Lookup("max_stack_buf_size"); f == nil {
		return fmt.Errorf("max_stack_buf_size is not a flag")
	}
	maxStackBufSizeSet := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == "max_stack_buf_size" {
			maxStackBufSizeSet = true
		}
	})
	if v, ok := flags["max_stack_buf_size"]; ok && !maxStackBufSizeSet {
		return fmt.Errorf("max_stack_buf_size unexpectedly set to %v", v)
	}
	return nil
}

func TestFlags(t *testing.T) {
	sh := modules.NewShell()
	defer sh.Cleanup(nil, nil)
	h, err := sh.Start("child", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if err = h.Shutdown(nil, os.Stderr); err != nil {
		t.Errorf("unexpected error: %v", err)
	}
}
<commit_msg>lib/modules: change the NewShell factory to take a security.Principal arg.<commit_after>package vlog_test

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"testing"

	"veyron.io/veyron/veyron/lib/modules"
	"veyron.io/veyron/veyron/lib/testutil"

	"veyron.io/veyron/veyron2/vlog"
)

func TestHelperProcess(t *testing.T) {
	modules.DispatchInTest()
}

func init() {
	testutil.Init()
	modules.RegisterChild("child", "", child)
}

func child(stdin io.Reader, stdout, stderr io.Writer, env map[string]string, args ...string) error {
	tmp := filepath.Join(os.TempDir(), "foo")
	flag.Set("log_dir", tmp)
	flag.Set("vmodule", "foo=2")
	flags := vlog.Log.ExplicitlySetFlags()
	if v, ok := flags["log_dir"]; !ok || v != tmp {
		return fmt.Errorf("log_dir was supposed to be %v", tmp)
	}
	if v, ok := flags["vmodule"]; !ok || v != "foo=2" {
		return fmt.Errorf("vmodule was supposed to be foo=2")
	}
	if f := flag.Lookup("max_stack_buf_size"); f == nil {
		return fmt.Errorf("max_stack_buf_size is not a flag")
	}
	maxStackBufSizeSet := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == "max_stack_buf_size" {
			maxStackBufSizeSet = true
		}
	})
	if v, ok := flags["max_stack_buf_size"]; ok && !maxStackBufSizeSet {
		return fmt.Errorf("max_stack_buf_size unexpectedly set to %v", v)
	}
	return nil
}

func TestFlags(t *testing.T) {
	sh, err := modules.NewShell(nil)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}
	defer sh.Cleanup(nil, nil)
	h, err := sh.Start("child", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if err = h.Shutdown(nil, os.Stderr); err != nil {
		t.Errorf("unexpected error: %v", err)
	}
}
<|endoftext|>


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

package vm

import (
	"strings"
	"testing"

	"github.com/go-test/deep"
)

var testCodeGenPrograms = []struct {
	name   string
	source string
	prog   []instr // expected bytecode
}{
	// Composite literals require too many explicit conversions.
	{"simple line counter",
		"counter line_count\n/$/ { line_count++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"count a",
		"counter a_count\n/a$/ { a_count++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"strptime and capref",
		"counter foo\n" +
			"/(.*)/ { strptime($1, \"2006-01-02T15:04:05\")\n" +
			"foo++\n}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 11},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{str, 0},
			instr{strptime, 2},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"strptime and named capref",
		"counter foo\n" +
			"/(?P<date>.*)/ { strptime($date, \"2006-01-02T15:04:05\")\n" +
			"foo++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 11},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 1},
			instr{str, 0},
			instr{strptime, 2},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"inc by and set",
		"counter foo\ncounter bar\n" +
			"/([0-9]+)/ {\n" +
			"foo += $1\n" +
			"bar = $1\n" +
			"}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 17},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{mload, 0},
			instr{dload, 0},
			instr{push, 0},
			instr{capref, 0},
			instr{iadd, nil},
			instr{iset, nil},
			instr{mload, 1},
			instr{dload, 0},
			instr{push, 0},
			instr{capref, 0},
			instr{iset, nil},
			instr{setmatched, true}}},
	{"cond expr gt",
		"counter foo\n" +
			"1 > 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr lt",
		"counter foo\n" +
			"1 < 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, -1},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr eq",
		"counter foo\n" +
			"1 == 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 0},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr le",
		"counter foo\n" +
			"1 <= 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr ge",
		"counter foo\n" +
			"1 >= 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, -1},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr ne",
		"counter foo\n" +
			"1 != 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 0},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"nested cond",
		"counter foo\n" +
			"/(.*)/ {\n" +
			"  $1 <= 1 {\n" +
			"    foo++\n" +
			"  }\n" +
			"}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 14},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{push, 1},
			instr{cmp, 1},
			instr{jm, 13},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true},
			instr{setmatched, true}}},
	{"deco",
		"counter foo\n" +
			"counter bar\n" +
			"def fooWrap {\n" +
			"  /.*/ {\n" +
			"    foo++\n" +
			"    next\n" +
			"  }\n" +
			"}\n" +
			"" +
			"@fooWrap { bar++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 10},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{mload, 1},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"length",
		"len(\"foo\") > 0 {\n" +
			"}\n",
		[]instr{
			instr{str, 0},
			instr{length, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{setmatched, true}}},
	{"bitwise", `
1 & 7 ^ 15 | 8
~ 16 << 2
1 >> 20
`,
		[]instr{
			instr{push, 1},
			instr{push, 7},
			instr{and, nil},
			instr{push, 15},
			instr{xor, nil},
			instr{push, 8},
			instr{or, nil},
			instr{push, 16},
			instr{not, nil},
			instr{push, 2},
			instr{shl, nil},
			instr{push, 1},
			instr{push, 20},
			instr{shr, nil}}},
	{"pow", `
/(\d+) (\d+)/ {
$1 ** $2
}
`,
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{push, 0},
			instr{capref, 1},
			instr{ipow, nil},
			instr{setmatched, true}}},
	{"indexed expr", `
counter a by b
a["string"]++
`,
		[]instr{
			instr{str, 0},
			instr{mload, 0},
			instr{dload, 1},
			instr{inc, nil}}},
	{"strtol", `
strtol("deadbeef", 16)
`,
		[]instr{
			instr{str, 0},
			instr{push, 16},
			instr{strtol, 2}}},
	{"float", `
20.0
`,
		[]instr{
			instr{push, 20.0}}},
	{"otherwise", `
counter a
otherwise {
	a++
}
`,
		[]instr{
			instr{otherwise, nil},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond else",
		`counter foo
counter bar
1 > 0 {
  foo++
} else {
  bar++
}`,
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 10},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true},
			instr{jmp, 13},
			instr{mload, 1},
			instr{dload, 0},
			instr{inc, nil},
		},
	},
	{"mod",
		`
3 % 1
`,
		[]instr{
			instr{push, 3},
			instr{push, 1},
			instr{imod, nil},
		},
	},
	{"del", `
counter a by b
del a["string"]
`,
		[]instr{
			instr{str, 0},
			instr{mload, 0},
			instr{del, 1}},
	},
}

func TestCodegen(t *testing.T) {
	for _, tc := range testCodeGenPrograms {
		t.Logf("Starting %s", tc.name)
		ast, err := Parse(tc.name, strings.NewReader(tc.source))
		if err != nil {
			t.Fatalf("Unexpected parse failure in %q: %s", tc.name, err)
		}
		err = Check(ast)
		if err != nil {
			t.Fatalf("Unexpected check failure in %q: %s", tc.name, err)
		}
		obj, err := CodeGen(tc.name, ast)
		if err != nil {
			t.Errorf("Compile errors for %q:\n%s", tc.name, err)
			s := Sexp{}
			s.emitTypes = true
			t.Logf("AST:\n%s", s.Dump(ast))
			continue
		}
		if diff := deep.Equal(tc.prog, obj.prog); diff != nil {
			t.Errorf("%s: VM prog doesn't match.\n%s", tc.name, diff)
		}
	}
}
<commit_msg>Add a unit test for codegen on type of capref inference.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

import (
	"strings"
	"testing"

	"github.com/go-test/deep"
)

var testCodeGenPrograms = []struct {
	name   string
	source string
	prog   []instr // expected bytecode
}{
	// Composite literals require too many explicit conversions.
	{"simple line counter",
		"counter line_count\n/$/ { line_count++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"count a",
		"counter a_count\n/a$/ { a_count++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"strptime and capref",
		"counter foo\n" +
			"/(.*)/ { strptime($1, \"2006-01-02T15:04:05\")\n" +
			"foo++\n}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 11},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{str, 0},
			instr{strptime, 2},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"strptime and named capref",
		"counter foo\n" +
			"/(?P<date>.*)/ { strptime($date, \"2006-01-02T15:04:05\")\n" +
			"foo++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 11},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 1},
			instr{str, 0},
			instr{strptime, 2},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"inc by and set",
		"counter foo\ncounter bar\n" +
			"/([0-9]+)/ {\n" +
			"foo += $1\n" +
			"bar = $1\n" +
			"}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 17},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{mload, 0},
			instr{dload, 0},
			instr{push, 0},
			instr{capref, 0},
			instr{iadd, nil},
			instr{iset, nil},
			instr{mload, 1},
			instr{dload, 0},
			instr{push, 0},
			instr{capref, 0},
			instr{iset, nil},
			instr{setmatched, true}}},
	{"cond expr gt",
		"counter foo\n" +
			"1 > 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr lt",
		"counter foo\n" +
			"1 < 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, -1},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr eq",
		"counter foo\n" +
			"1 == 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 0},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr le",
		"counter foo\n" +
			"1 <= 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr ge",
		"counter foo\n" +
			"1 >= 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, -1},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond expr ne",
		"counter foo\n" +
			"1 != 0 {\n" +
			"  foo++\n" +
			"}\n",
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 0},
			instr{jm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"nested cond",
		"counter foo\n" +
			"/(.*)/ {\n" +
			"  $1 <= 1 {\n" +
			"    foo++\n" +
			"  }\n" +
			"}\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 14},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{push, 1},
			instr{cmp, 1},
			instr{jm, 13},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true},
			instr{setmatched, true}}},
	{"deco",
		"counter foo\n" +
			"counter bar\n" +
			"def fooWrap {\n" +
			"  /.*/ {\n" +
			"    foo++\n" +
			"    next\n" +
			"  }\n" +
			"}\n" +
			"" +
			"@fooWrap { bar++\n }\n",
		[]instr{
			instr{match, 0},
			instr{jnm, 10},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{mload, 1},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"length",
		"len(\"foo\") > 0 {\n" +
			"}\n",
		[]instr{
			instr{str, 0},
			instr{length, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{setmatched, true}}},
	{"bitwise", `
1 & 7 ^ 15 | 8
~ 16 << 2
1 >> 20
`,
		[]instr{
			instr{push, 1},
			instr{push, 7},
			instr{and, nil},
			instr{push, 15},
			instr{xor, nil},
			instr{push, 8},
			instr{or, nil},
			instr{push, 16},
			instr{not, nil},
			instr{push, 2},
			instr{shl, nil},
			instr{push, 1},
			instr{push, 20},
			instr{shr, nil}}},
	{"pow", `
/(\d+) (\d+)/ {
$1 ** $2
}
`,
		[]instr{
			instr{match, 0},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{push, 0},
			instr{capref, 0},
			instr{push, 0},
			instr{capref, 1},
			instr{ipow, nil},
			instr{setmatched, true}}},
	{"indexed expr", `
counter a by b
a["string"]++
`,
		[]instr{
			instr{str, 0},
			instr{mload, 0},
			instr{dload, 1},
			instr{inc, nil}}},
	{"strtol", `
strtol("deadbeef", 16)
`,
		[]instr{
			instr{str, 0},
			instr{push, 16},
			instr{strtol, 2}}},
	{"float", `
20.0
`,
		[]instr{
			instr{push, 20.0}}},
	{"otherwise", `
counter a
otherwise {
	a++
}
`,
		[]instr{
			instr{otherwise, nil},
			instr{jnm, 7},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true}}},
	{"cond else",
		`counter foo
counter bar
1 > 0 {
  foo++
} else {
  bar++
}`,
		[]instr{
			instr{push, 1},
			instr{push, 0},
			instr{cmp, 1},
			instr{jnm, 10},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{inc, nil},
			instr{setmatched, true},
			instr{jmp, 13},
			instr{mload, 1},
			instr{dload, 0},
			instr{inc, nil},
		},
	},
	{"mod",
		`
3 % 1
`,
		[]instr{
			instr{push, 3},
			instr{push, 1},
			instr{imod, nil},
		},
	},
	{"del", `
counter a by b
del a["string"]
`,
		[]instr{
			instr{str, 0},
			instr{mload, 0},
			instr{del, 1}},
	},
	{"types", `
gauge i
gauge f
/(\d+)/ {
 i = $1
}
/(\d+\.\d+)/ {
 f = $1
}
`,
		[]instr{
			instr{match, 0},
			instr{jnm, 9},
			instr{setmatched, false},
			instr{mload, 0},
			instr{dload, 0},
			instr{push, 0},
			instr{capref, 1},
			instr{iset, nil},
			instr{setmatched, true},
			instr{match, 0},
			instr{jnm, 18},
			instr{setmatched, false},
			instr{mload, 1},
			instr{dload, 0},
			instr{push, 1},
			instr{capref, 1},
			instr{fset, nil},
			instr{setmatched, true},
		},
	},
}

func TestCodegen(t *testing.T) {
	for _, tc := range testCodeGenPrograms {
		t.Logf("Starting %s", tc.name)
		ast, err := Parse(tc.name, strings.NewReader(tc.source))
		if err != nil {
			t.Fatalf("Unexpected parse failure in %q: %s", tc.name, err)
		}
		err = Check(ast)
		if err != nil {
			t.Fatalf("Unexpected check failure in %q: %s", tc.name, err)
		}
		obj, err := CodeGen(tc.name, ast)
		if err != nil {
			t.Errorf("Compile errors for %q:\n%s", tc.name, err)
			s := Sexp{}
			s.emitTypes = true
			t.Logf("AST:\n%s", s.Dump(ast))
			continue
		}
		if diff := deep.Equal(tc.prog, obj.prog); diff != nil {
			t.Errorf("%s: VM prog doesn't match.\n%s", tc.name, diff)
			t.Logf("Expected:\n%s\nReceived:\n%s", tc.prog, obj.prog)
		}
	}
}
<|endoftext|>