<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package leveldb

import (
	"time"

	"github.com/syndtr/goleveldb/leveldb/memdb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/util"
)

func (db *DB) writeJournal(b *Batch) error {
	w, err := db.journal.Next()
	if err != nil {
		return err
	}
	if _, err := w.Write(b.encode()); err != nil {
		return err
	}
	if err := db.journal.Flush(); err != nil {
		return err
	}
	if b.sync {
		return db.journalWriter.Sync()
	}
	return nil
}

func (db *DB) jWriter() {
	defer db.closeW.Done()
	for {
		select {
		case b := <-db.journalC:
			if b != nil {
				db.journalAckC <- db.writeJournal(b)
			}
		case _, _ = <-db.closeC:
			return
		}
	}
}

func (db *DB) rotateMem(n int) (mem *memdb.DB, err error) {
	// Wait for pending memdb compaction.
	err = db.compSendIdle(db.mcompCmdC)
	if err != nil {
		return
	}

	// Create new memdb and journal.
	mem, err = db.newMem(n)
	if err != nil {
		return
	}

	// Schedule memdb compaction.
	db.compTrigger(db.mcompTriggerC)
	return
}

func (db *DB) flush(n int) (mem *memdb.DB, nn int, err error) {
	delayed := false
	flush := func() bool {
		v := db.s.version()
		defer v.release()
		mem = db.getEffectiveMem()
		nn = mem.Free()
		switch {
		case v.tLen(0) >= kL0_SlowdownWritesTrigger && !delayed:
			delayed = true
			time.Sleep(time.Millisecond)
		case nn >= n:
			return false
		case v.tLen(0) >= kL0_StopWritesTrigger:
			delayed = true
			err = db.compSendIdle(db.tcompCmdC)
			if err != nil {
				return false
			}
		default:
			// Allow memdb to grow if it has no entry.
			if mem.Len() == 0 {
				nn = n
				return false
			}
			mem, err = db.rotateMem(n)
			nn = mem.Free()
			return false
		}
		return true
	}
	start := time.Now()
	for flush() {
	}
	if delayed {
		db.logf("db@write delayed T·%v", time.Since(start))
	}
	return
}

// Write apply the given batch to the DB. The batch will be applied
// sequentially.
//
// It is safe to modify the contents of the arguments after Write returns.
func (db *DB) Write(b *Batch, wo *opt.WriteOptions) (err error) {
	err = db.ok()
	if err != nil || b == nil || b.len() == 0 {
		return
	}

	b.init(wo.GetSync())

	// The write happen synchronously.
retry:
	select {
	case db.writeC <- b:
		if <-db.writeMergedC {
			return <-db.writeAckC
		}
		goto retry
	case db.writeLockC <- struct{}{}:
	case _, _ = <-db.closeC:
		return ErrClosed
	}

	merged := 0
	defer func() {
		<-db.writeLockC
		for i := 0; i < merged; i++ {
			db.writeAckC <- err
		}
	}()

	mem, memFree, err := db.flush(b.size())
	if err != nil {
		return
	}

	// Calculate maximum size of the batch.
	m := 1 << 20
	if x := b.size(); x <= 128<<10 {
		m = x + (128 << 10)
	}
	m = minInt(m, memFree)

	// Merge with other batch.
drain:
	for b.size() < m && !b.sync {
		select {
		case nb := <-db.writeC:
			if b.size()+nb.size() <= m {
				b.append(nb)
				db.writeMergedC <- true
				merged++
			} else {
				db.writeMergedC <- false
				break drain
			}
		default:
			break drain
		}
	}

	// Set batch first seq number relative from last seq.
	b.seq = db.seq + 1

	// Write journal concurrently if it is large enough.
	if b.size() >= (128 << 10) {
		// Push the write batch to the journal writer
		select {
		case _, _ = <-db.closeC:
			err = ErrClosed
			return
		case db.journalC <- b:
			// Write into memdb
			b.memReplay(mem)
		}
		// Wait for journal writer
		select {
		case _, _ = <-db.closeC:
			err = ErrClosed
			return
		case err = <-db.journalAckC:
			if err != nil {
				// Revert memdb if error detected
				b.revertMemReplay(mem)
				return
			}
		}
	} else {
		err = db.writeJournal(b)
		if err != nil {
			return
		}
		b.memReplay(mem)
	}

	// Set last seq number.
	db.addSeq(uint64(b.len()))

	if b.size() >= memFree {
		db.rotateMem(0)
	}
	return
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (db *DB) Put(key, value []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Put(key, value)
	return db.Write(b, wo)
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (db *DB) Delete(key []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Delete(key)
	return db.Write(b, wo)
}

func isMemOverlaps(icmp *iComparer, mem *memdb.DB, min, max []byte) bool {
	iter := mem.NewIterator(nil)
	defer iter.Release()
	return (max == nil || (iter.First() && icmp.uCompare(max, iKey(iter.Key()).ukey()) >= 0)) &&
		(min == nil || (iter.Last() && icmp.uCompare(min, iKey(iter.Key()).ukey()) <= 0))
}

// CompactRange compacts the underlying DB for the given key range.
// In particular, deleted and overwritten versions are discarded,
// and the data is rearranged to reduce the cost of operations
// needed to access the data. This operation should typically only
// be invoked by users who understand the underlying implementation.
//
// A nil Range.Start is treated as a key before all keys in the DB.
// And a nil Range.Limit is treated as a key after all keys in the DB.
// Therefore if both is nil then it will compact entire DB.
func (db *DB) CompactRange(r util.Range) error {
	if err := db.ok(); err != nil {
		return err
	}

	select {
	case db.writeLockC <- struct{}{}:
	case _, _ = <-db.closeC:
		return ErrClosed
	}

	// Check for overlaps in memdb.
	mem := db.getEffectiveMem()
	if isMemOverlaps(db.s.icmp, mem, r.Start, r.Limit) {
		// Memdb compaction.
		if _, err := db.rotateMem(0); err != nil {
			<-db.writeLockC
			return err
		}
		<-db.writeLockC
		if err := db.compSendIdle(db.mcompCmdC); err != nil {
			return err
		}
	} else {
		<-db.writeLockC
	}

	// Table compaction.
	return db.compSendRange(db.tcompCmdC, -1, r.Start, r.Limit)
}
<commit_msg>leveldb: Fix forgotten error check #63<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package leveldb

import (
	"time"

	"github.com/syndtr/goleveldb/leveldb/memdb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/util"
)

func (db *DB) writeJournal(b *Batch) error {
	w, err := db.journal.Next()
	if err != nil {
		return err
	}
	if _, err := w.Write(b.encode()); err != nil {
		return err
	}
	if err := db.journal.Flush(); err != nil {
		return err
	}
	if b.sync {
		return db.journalWriter.Sync()
	}
	return nil
}

func (db *DB) jWriter() {
	defer db.closeW.Done()
	for {
		select {
		case b := <-db.journalC:
			if b != nil {
				db.journalAckC <- db.writeJournal(b)
			}
		case _, _ = <-db.closeC:
			return
		}
	}
}

func (db *DB) rotateMem(n int) (mem *memdb.DB, err error) {
	// Wait for pending memdb compaction.
	err = db.compSendIdle(db.mcompCmdC)
	if err != nil {
		return
	}

	// Create new memdb and journal.
	mem, err = db.newMem(n)
	if err != nil {
		return
	}

	// Schedule memdb compaction.
	db.compTrigger(db.mcompTriggerC)
	return
}

func (db *DB) flush(n int) (mem *memdb.DB, nn int, err error) {
	delayed := false
	flush := func() bool {
		v := db.s.version()
		defer v.release()
		mem = db.getEffectiveMem()
		nn = mem.Free()
		switch {
		case v.tLen(0) >= kL0_SlowdownWritesTrigger && !delayed:
			delayed = true
			time.Sleep(time.Millisecond)
		case nn >= n:
			return false
		case v.tLen(0) >= kL0_StopWritesTrigger:
			delayed = true
			err = db.compSendIdle(db.tcompCmdC)
			if err != nil {
				return false
			}
		default:
			// Allow memdb to grow if it has no entry.
			if mem.Len() == 0 {
				nn = n
			} else {
				mem, err = db.rotateMem(n)
				if err == nil {
					nn = mem.Free()
				} else {
					nn = 0
				}
			}
			return false
		}
		return true
	}
	start := time.Now()
	for flush() {
	}
	if delayed {
		db.logf("db@write delayed T·%v", time.Since(start))
	}
	return
}

// Write apply the given batch to the DB. The batch will be applied
// sequentially.
//
// It is safe to modify the contents of the arguments after Write returns.
func (db *DB) Write(b *Batch, wo *opt.WriteOptions) (err error) {
	err = db.ok()
	if err != nil || b == nil || b.len() == 0 {
		return
	}

	b.init(wo.GetSync())

	// The write happen synchronously.
retry:
	select {
	case db.writeC <- b:
		if <-db.writeMergedC {
			return <-db.writeAckC
		}
		goto retry
	case db.writeLockC <- struct{}{}:
	case _, _ = <-db.closeC:
		return ErrClosed
	}

	merged := 0
	defer func() {
		<-db.writeLockC
		for i := 0; i < merged; i++ {
			db.writeAckC <- err
		}
	}()

	mem, memFree, err := db.flush(b.size())
	if err != nil {
		return
	}

	// Calculate maximum size of the batch.
	m := 1 << 20
	if x := b.size(); x <= 128<<10 {
		m = x + (128 << 10)
	}
	m = minInt(m, memFree)

	// Merge with other batch.
drain:
	for b.size() < m && !b.sync {
		select {
		case nb := <-db.writeC:
			if b.size()+nb.size() <= m {
				b.append(nb)
				db.writeMergedC <- true
				merged++
			} else {
				db.writeMergedC <- false
				break drain
			}
		default:
			break drain
		}
	}

	// Set batch first seq number relative from last seq.
	b.seq = db.seq + 1

	// Write journal concurrently if it is large enough.
	if b.size() >= (128 << 10) {
		// Push the write batch to the journal writer
		select {
		case _, _ = <-db.closeC:
			err = ErrClosed
			return
		case db.journalC <- b:
			// Write into memdb
			b.memReplay(mem)
		}
		// Wait for journal writer
		select {
		case _, _ = <-db.closeC:
			err = ErrClosed
			return
		case err = <-db.journalAckC:
			if err != nil {
				// Revert memdb if error detected
				b.revertMemReplay(mem)
				return
			}
		}
	} else {
		err = db.writeJournal(b)
		if err != nil {
			return
		}
		b.memReplay(mem)
	}

	// Set last seq number.
	db.addSeq(uint64(b.len()))

	if b.size() >= memFree {
		db.rotateMem(0)
	}
	return
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (db *DB) Put(key, value []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Put(key, value)
	return db.Write(b, wo)
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (db *DB) Delete(key []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Delete(key)
	return db.Write(b, wo)
}

func isMemOverlaps(icmp *iComparer, mem *memdb.DB, min, max []byte) bool {
	iter := mem.NewIterator(nil)
	defer iter.Release()
	return (max == nil || (iter.First() && icmp.uCompare(max, iKey(iter.Key()).ukey()) >= 0)) &&
		(min == nil || (iter.Last() && icmp.uCompare(min, iKey(iter.Key()).ukey()) <= 0))
}

// CompactRange compacts the underlying DB for the given key range.
// In particular, deleted and overwritten versions are discarded,
// and the data is rearranged to reduce the cost of operations
// needed to access the data. This operation should typically only
// be invoked by users who understand the underlying implementation.
//
// A nil Range.Start is treated as a key before all keys in the DB.
// And a nil Range.Limit is treated as a key after all keys in the DB.
// Therefore if both is nil then it will compact entire DB.
func (db *DB) CompactRange(r util.Range) error {
	if err := db.ok(); err != nil {
		return err
	}

	select {
	case db.writeLockC <- struct{}{}:
	case _, _ = <-db.closeC:
		return ErrClosed
	}

	// Check for overlaps in memdb.
	mem := db.getEffectiveMem()
	if isMemOverlaps(db.s.icmp, mem, r.Start, r.Limit) {
		// Memdb compaction.
		if _, err := db.rotateMem(0); err != nil {
			<-db.writeLockC
			return err
		}
		<-db.writeLockC
		if err := db.compSendIdle(db.mcompCmdC); err != nil {
			return err
		}
	} else {
		<-db.writeLockC
	}

	// Table compaction.
	return db.compSendRange(db.tcompCmdC, -1, r.Start, r.Limit)
}
<|endoftext|>


<commit_before>package sqlite

import (
	"reflect"
	"testing"

	"../godata"
)

const SQLITE3_TEST_TABLE_NAME = "FOO"
const SQLITE3_TEST_CREATE_SCHEMA = "CREATE TABLE FOO (NAME TEXT, ADDRESS TEXT, ID INTEGER, CREATED_TIMESTAMP DEFAULT CURRENT_TIMESTAMP);"
const SQLITE3_TEST_DROP_SCHEMA = "DROP TABLE FOO"
const SQLITE3_TEST_INSERT = "INSERT INTO FOO (NAME,ADDRESS) VALUES (?,?);"
const SQLITE3_TEST_SELECT = "SELECT NAME,ADDRESS FROM FOO;"
const SQLITE3_TEST_UPDATE = "UPDATE FOO SET ADDRESS = ? WHERE NAME = ?"
const SQLITE3_TEST_CREATE_SCHEMA_FAIL = "CREATE TABLE FOO (NAME TEXT, ADDRESS TEXT, ID2 INTEGER, CREATED_TIMESTAMP DEFAULT CURRENT_TIMESTAMP);"

func TestNewConnection(t *testing.T) {
	_, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
}

func TestClose(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
	err = conn.Close()
	if err != nil {
		t.Error(conn)
	}

}

func TestCheckTableExists(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
	exists, err := conn.CheckTableExists(SQLITE3_TEST_TABLE_NAME)
	if err != nil {
		t.Error(err)
	}

	if exists != false {
		t.Error("error => table " + SQLITE3_TEST_TABLE_NAME + "should not exist")
	}

	err = conn.Close()
	if err != nil {
		t.Error(conn)
	}
}

func TestInitTable(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	err = conn.Close()

	if err != nil {
		t.Error(err)
	}

}

func TestInsert(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "austin")

	if err != nil {
		t.Error(err)
	}

	if n != 1 {
		t.Error("Row was not inserted properly ; n = ", n)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestUpdate(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "austin")

	if err != nil {
		t.Error(err)
	}

	if n != 1 {
		t.Error("Row was not inserted properly ; n = ", n)
	}

	n, err = conn.Insert(SQLITE3_TEST_INSERT, "bmw", "austin")

	if err != nil {
		t.Error(err)
	}

	if n != 2 {
		t.Error("Row was not inserted properly ; n = ", n)
	}

	err = conn.Update(SQLITE3_TEST_UPDATE, "boston", "parag")

	if err != nil {
		t.Error(err)
	}

	_, err = conn.Select(SQLITE3_TEST_SELECT)

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestSelect(t *testing.T) {

	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "austinYO")

	if n != 1 {
		t.Error("statement is not inserted n => ", n)
	}

	if err != nil {
		t.Error(err)
	}

	_, err = conn.Select(SQLITE3_TEST_SELECT)

	if err != nil {
		t.Error(err)
	}

	//_:= goselect.GetColumns()
	//ows := goselect.GetRows()

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestMetaGoTable(t *testing.T) {

	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	gmeta, err := conn.MetaGoTable(SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	err = conn.Close()

	if err != nil {
		t.Error(err)
	}

	gtest, err := godata.NewMetaTableFromCreateStatement(SQLITE3_TEST_CREATE_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	if !reflect.DeepEqual(gmeta, gtest) {
		t.Error("gmeta and gtest are not the same")
		t.Error("gmeta column values are ", gmeta.GetColumns())
		t.Error("gtest column values are ", gtest.GetColumns())
	}

}
<commit_msg>Added to the TestUpdated test case.  Will now query the database and make sure the results of the update are correct based on the returned data.<commit_after>package sqlite

import (
	"reflect"
	"testing"

	"../godata"
)

const SQLITE3_TEST_TABLE_NAME = "FOO"
const SQLITE3_TEST_CREATE_SCHEMA = "CREATE TABLE FOO (NAME TEXT, ADDRESS TEXT, ID INTEGER, CREATED_TIMESTAMP DEFAULT CURRENT_TIMESTAMP);"
const SQLITE3_TEST_DROP_SCHEMA = "DROP TABLE FOO"
const SQLITE3_TEST_INSERT = "INSERT INTO FOO (NAME,ADDRESS) VALUES (?,?);"
const SQLITE3_TEST_SELECT = "SELECT NAME,ADDRESS FROM FOO;"
const SQLITE3_TEST_UPDATE = "UPDATE FOO SET ADDRESS = ? WHERE NAME = ?"
const SQLITE3_TEST_CREATE_SCHEMA_FAIL = "CREATE TABLE FOO (NAME TEXT, ADDRESS TEXT, ID2 INTEGER, CREATED_TIMESTAMP DEFAULT CURRENT_TIMESTAMP);"

func TestNewConnection(t *testing.T) {
	_, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
}

func TestClose(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
	err = conn.Close()
	if err != nil {
		t.Error(conn)
	}

}

func TestCheckTableExists(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}
	exists, err := conn.CheckTableExists(SQLITE3_TEST_TABLE_NAME)
	if err != nil {
		t.Error(err)
	}

	if exists != false {
		t.Error("error => table " + SQLITE3_TEST_TABLE_NAME + "should not exist")
	}

	err = conn.Close()
	if err != nil {
		t.Error(conn)
	}
}

func TestInitTable(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	err = conn.Close()

	if err != nil {
		t.Error(err)
	}

}

func TestInsert(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "austin")

	if err != nil {
		t.Error(err)
	}

	if n != 1 {
		t.Error("Row was not inserted properly ; n = ", n)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestUpdate(t *testing.T) {
	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "boston")

	if err != nil {
		t.Error(err)
	}

	if n != 1 {
		t.Error("Row was not inserted properly ; n = ", n)
	}

	gdata, err := conn.Select(SQLITE3_TEST_SELECT)

	if gdata.GetRows()[0]["NAME"] != "parag" {
		t.Error("Name not inserted correctly")
	}

	if gdata.GetRows()[0]["ADDRESS"] != "boston" {
		t.Error("Name not inserted correctly")
	}

	err = conn.Update(SQLITE3_TEST_UPDATE, "austin", "parag")

	if err != nil {
		t.Error(err)
	}

	gdata, err = conn.Select(SQLITE3_TEST_SELECT)

	if gdata.GetRows()[0]["ADDRESS"] != "austin" {
		t.Error("Name not inserted correctly")
	}

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestSelect(t *testing.T) {

	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	defer conn.Close()

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	n, err := conn.Insert(SQLITE3_TEST_INSERT, "parag", "austinYO")

	if n != 1 {
		t.Error("statement is not inserted n => ", n)
	}

	if err != nil {
		t.Error(err)
	}

	_, err = conn.Select(SQLITE3_TEST_SELECT)

	if err != nil {
		t.Error(err)
	}

	//_:= goselect.GetColumns()
	//ows := goselect.GetRows()

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

}

func TestMetaGoTable(t *testing.T) {

	conn, err := NewConnection()
	if err != nil {
		t.Error(err)
	}

	err = conn.InitTable(SQLITE3_TEST_CREATE_SCHEMA, SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	gmeta, err := conn.MetaGoTable(SQLITE3_TEST_TABLE_NAME)

	if err != nil {
		t.Error(err)
	}

	_, err = conn.db.Exec(SQLITE3_TEST_DROP_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	err = conn.Close()

	if err != nil {
		t.Error(err)
	}

	gtest, err := godata.NewMetaTableFromCreateStatement(SQLITE3_TEST_CREATE_SCHEMA)

	if err != nil {
		t.Error(err)
	}

	if !reflect.DeepEqual(gmeta, gtest) {
		t.Error("gmeta and gtest are not the same")
		t.Error("gmeta column values are ", gmeta.GetColumns())
		t.Error("gtest column values are ", gtest.GetColumns())
	}

}
<|endoftext|>


<commit_before>/*
 * This file is part of the kubevirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package isolation

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

/*
 ATTENTION: Rerun code generators when interface signatures are modified.
*/

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"kubevirt.io/kubevirt/pkg/unsafepath"

	mount "github.com/moby/sys/mountinfo"

	"kubevirt.io/kubevirt/pkg/safepath"

	"kubevirt.io/client-go/log"

	"kubevirt.io/kubevirt/pkg/util"
)

// IsolationResult is the result of a successful PodIsolationDetector.Detect
type IsolationResult interface {
	// process ID
	Pid() int
	// parent process ID
	PPid() int
	// full path to the process namespace
	PIDNamespace() string
	// full path to the process root mount
	MountRoot() (*safepath.Path, error)
	// full path to the mount namespace
	MountNamespace() string
	// mounts for the process
	Mounts(mount.FilterFunc) ([]*mount.Info, error)
}

type RealIsolationResult struct {
	pid  int
	ppid int
}

func NewIsolationResult(pid, ppid int) IsolationResult {
	return &RealIsolationResult{pid: pid, ppid: ppid}
}

func (r *RealIsolationResult) PIDNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/pid", r.pid)
}

func (r *RealIsolationResult) MountNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/mnt", r.pid)
}

// IsMounted checks if the given path is a mount point or not.
func IsMounted(mountPoint *safepath.Path) (isMounted bool, err error) {
	// Ensure that the path is still a valid absolute path without symlinks
	f, err := safepath.OpenAtNoFollow(mountPoint)
	if err != nil {
		// treat os.IsNotExist() as error too
		// since the inherent property of a safepath.Path is that the path must
		// have existed at the point of object creation
		return false, err
	}
	defer f.Close()
	if mountPoint.IsRoot() {
		// mount.Mounted has purely string matching based special logic on how to treat "/".
		// Emulating this for safepath here without ever having to call an unsafe method on our
		// safepath.
		return true, nil
	} else {
		// TODO: Unsafe full path is required, and not a fd, since otherwise mount table lookups and such would not work.
		return mount.Mounted(unsafepath.UnsafeAbsolute(mountPoint.Raw()))
	}
}

// AreMounted checks if given paths are mounted by calling IsMounted.
// If error occurs, the first error is returned.
func (r *RealIsolationResult) AreMounted(mountPoints ...*safepath.Path) (isMounted bool, err error) {
	for _, mountPoint := range mountPoints {
		isMounted, err = IsMounted(mountPoint)
		if !isMounted || err != nil {
			return
		}
	}

	return true, nil
}

// IsBlockDevice checks if the given path is a block device or not.
func IsBlockDevice(path *safepath.Path) (bool, error) {
	fileInfo, err := safepath.StatAtNoFollow(path)
	if err != nil {
		return false, fmt.Errorf("error checking for block device: %v", err)
	}
	if fileInfo.IsDir() || (fileInfo.Mode()&os.ModeDevice) == 0 {
		return false, nil
	}
	return true, nil
}

func (r *RealIsolationResult) MountRoot() (*safepath.Path, error) {
	return safepath.JoinAndResolveWithRelativeRoot(fmt.Sprintf("/proc/%d/root", r.pid))
}

func (r *RealIsolationResult) MountRootRelative(relativePath string) (*safepath.Path, error) {
	mountRoot, err := r.MountRoot()
	if err != nil {
		return nil, err
	}
	return mountRoot.AppendAndResolveWithRelativeRoot(relativePath)
}

func (r *RealIsolationResult) Pid() int {
	return r.pid
}

func (r *RealIsolationResult) PPid() int {
	return r.ppid
}

func NodeIsolationResult() *RealIsolationResult {
	return &RealIsolationResult{
		pid: 1,
	}
}

// Mounts returns mounts for the given process based on the supplied filter
func (r *RealIsolationResult) Mounts(filter mount.FilterFunc) ([]*mount.Info, error) {
	in, err := os.Open(fmt.Sprintf("/proc/%d/mountinfo", r.pid))
	if err != nil {
		return nil, fmt.Errorf("could not open file mountinfo for %d: %v", r.pid, err)
	}
	defer util.CloseIOAndCheckErr(in, nil)
	return mount.GetMountsFromReader(in, filter)
}

// MountInfoRoot returns the mount information for the root mount point
func MountInfoRoot(r IsolationResult) (mountInfo *mount.Info, err error) {
	mounts, err := r.Mounts(mount.SingleEntryFilter("/"))
	if err != nil {
		return nil, fmt.Errorf("failed to process mountinfo for pid %d: %v", r.Pid(), err)
	}
	if len(mounts) <= 0 {
		return nil, fmt.Errorf("no root mount point entry found for pid %d", r.Pid())
	}
	return mounts[0], nil
}

// parentMountInfoFor takes the mountInfo record of a container (child) and
// attempts to locate a mountpoint containing it on the parent.
func parentMountInfoFor(parent IsolationResult, mountInfo *mount.Info) (*mount.Info, error) {
	mounts, err := parent.Mounts(func(m *mount.Info) (bool, bool) {
		return m.Major != mountInfo.Major || m.Minor != mountInfo.Minor ||
			!strings.HasPrefix(mountInfo.Root, m.Root), false
	})
	if err != nil {
		return nil, fmt.Errorf("failed to find mount for %v in the mount namespace of pid %d", mountInfo.Root, parent.Pid())
	}

	if len(mounts) <= 0 {
		return nil, fmt.Errorf("no mount containing %v found in the mount namespace of pid %d", mountInfo.Root, parent.Pid())
	} else if len(mounts) > 1 {
		log.Log.Infof("found %d possible mount point candidates for path %v", len(mounts), mountInfo.Root)
		sort.SliceStable(mounts, func(i, j int) bool {
			return len(mounts[i].Root) > len(mounts[j].Root)
		})
	}

	return mounts[0], nil
}

// ParentPathForRootMount takes a container (child) and composes a path to
// the root mount point in the context of the parent.
func ParentPathForRootMount(parent IsolationResult, child IsolationResult) (*safepath.Path, error) {
	childRootMountInfo, err := MountInfoRoot(child)
	if err != nil {
		return nil, err
	}
	parentMountInfo, err := parentMountInfoFor(parent, childRootMountInfo)
	if err != nil {
		return nil, err
	}
	parentMountRoot, err := parent.MountRoot()
	if err != nil {
		return nil, err
	}
	path := parentMountRoot
	path, err = path.AppendAndResolveWithRelativeRoot(parentMountInfo.Mountpoint)
	if err != nil {
		return nil, err
	}
	return path.AppendAndResolveWithRelativeRoot(strings.TrimPrefix(childRootMountInfo.Root, parentMountInfo.Root))
}

func SafeJoin(res IsolationResult, elems ...string) (*safepath.Path, error) {
	mountRoot, err := res.MountRoot()
	if err != nil {
		return nil, err
	}
	return mountRoot.AppendAndResolveWithRelativeRoot(elems...)
}
<commit_msg>Ensure any provided mount point that is empty is ignored when checking for mount state<commit_after>/*
 * This file is part of the kubevirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package isolation

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

/*
 ATTENTION: Rerun code generators when interface signatures are modified.
*/

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"kubevirt.io/kubevirt/pkg/unsafepath"

	mount "github.com/moby/sys/mountinfo"

	"kubevirt.io/kubevirt/pkg/safepath"

	"kubevirt.io/client-go/log"

	"kubevirt.io/kubevirt/pkg/util"
)

// IsolationResult is the result of a successful PodIsolationDetector.Detect
type IsolationResult interface {
	// process ID
	Pid() int
	// parent process ID
	PPid() int
	// full path to the process namespace
	PIDNamespace() string
	// full path to the process root mount
	MountRoot() (*safepath.Path, error)
	// full path to the mount namespace
	MountNamespace() string
	// mounts for the process
	Mounts(mount.FilterFunc) ([]*mount.Info, error)
}

type RealIsolationResult struct {
	pid  int
	ppid int
}

func NewIsolationResult(pid, ppid int) IsolationResult {
	return &RealIsolationResult{pid: pid, ppid: ppid}
}

func (r *RealIsolationResult) PIDNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/pid", r.pid)
}

func (r *RealIsolationResult) MountNamespace() string {
	return fmt.Sprintf("/proc/%d/ns/mnt", r.pid)
}

// IsMounted checks if the given path is a mount point or not.
func IsMounted(mountPoint *safepath.Path) (isMounted bool, err error) {
	// Ensure that the path is still a valid absolute path without symlinks
	f, err := safepath.OpenAtNoFollow(mountPoint)
	if err != nil {
		// treat os.IsNotExist() as error too
		// since the inherent property of a safepath.Path is that the path must
		// have existed at the point of object creation
		return false, err
	}
	defer f.Close()
	if mountPoint.IsRoot() {
		// mount.Mounted has purely string matching based special logic on how to treat "/".
		// Emulating this for safepath here without ever having to call an unsafe method on our
		// safepath.
		return true, nil
	} else {
		// TODO: Unsafe full path is required, and not a fd, since otherwise mount table lookups and such would not work.
		return mount.Mounted(unsafepath.UnsafeAbsolute(mountPoint.Raw()))
	}
}

// AreMounted checks if given paths are mounted by calling IsMounted.
// If error occurs, the first error is returned.
func (r *RealIsolationResult) AreMounted(mountPoints ...*safepath.Path) (isMounted bool, err error) {
	for _, mountPoint := range mountPoints {
		if mountPoint != nil {
			isMounted, err = IsMounted(mountPoint)
			if !isMounted || err != nil {
				return
			}
		}
	}

	return true, nil
}

// IsBlockDevice checks if the given path is a block device or not.
func IsBlockDevice(path *safepath.Path) (bool, error) {
	fileInfo, err := safepath.StatAtNoFollow(path)
	if err != nil {
		return false, fmt.Errorf("error checking for block device: %v", err)
	}
	if fileInfo.IsDir() || (fileInfo.Mode()&os.ModeDevice) == 0 {
		return false, nil
	}
	return true, nil
}

func (r *RealIsolationResult) MountRoot() (*safepath.Path, error) {
	return safepath.JoinAndResolveWithRelativeRoot(fmt.Sprintf("/proc/%d/root", r.pid))
}

func (r *RealIsolationResult) MountRootRelative(relativePath string) (*safepath.Path, error) {
	mountRoot, err := r.MountRoot()
	if err != nil {
		return nil, err
	}
	return mountRoot.AppendAndResolveWithRelativeRoot(relativePath)
}

func (r *RealIsolationResult) Pid() int {
	return r.pid
}

func (r *RealIsolationResult) PPid() int {
	return r.ppid
}

func NodeIsolationResult() *RealIsolationResult {
	return &RealIsolationResult{
		pid: 1,
	}
}

// Mounts returns mounts for the given process based on the supplied filter
func (r *RealIsolationResult) Mounts(filter mount.FilterFunc) ([]*mount.Info, error) {
	in, err := os.Open(fmt.Sprintf("/proc/%d/mountinfo", r.pid))
	if err != nil {
		return nil, fmt.Errorf("could not open file mountinfo for %d: %v", r.pid, err)
	}
	defer util.CloseIOAndCheckErr(in, nil)
	return mount.GetMountsFromReader(in, filter)
}

// MountInfoRoot returns the mount information for the root mount point
func MountInfoRoot(r IsolationResult) (mountInfo *mount.Info, err error) {
	mounts, err := r.Mounts(mount.SingleEntryFilter("/"))
	if err != nil {
		return nil, fmt.Errorf("failed to process mountinfo for pid %d: %v", r.Pid(), err)
	}
	if len(mounts) <= 0 {
		return nil, fmt.Errorf("no root mount point entry found for pid %d", r.Pid())
	}
	return mounts[0], nil
}

// parentMountInfoFor takes the mountInfo record of a container (child) and
// attempts to locate a mountpoint containing it on the parent.
func parentMountInfoFor(parent IsolationResult, mountInfo *mount.Info) (*mount.Info, error) {
	mounts, err := parent.Mounts(func(m *mount.Info) (bool, bool) {
		return m.Major != mountInfo.Major || m.Minor != mountInfo.Minor ||
			!strings.HasPrefix(mountInfo.Root, m.Root), false
	})
	if err != nil {
		return nil, fmt.Errorf("failed to find mount for %v in the mount namespace of pid %d", mountInfo.Root, parent.Pid())
	}

	if len(mounts) <= 0 {
		return nil, fmt.Errorf("no mount containing %v found in the mount namespace of pid %d", mountInfo.Root, parent.Pid())
	} else if len(mounts) > 1 {
		log.Log.Infof("found %d possible mount point candidates for path %v", len(mounts), mountInfo.Root)
		sort.SliceStable(mounts, func(i, j int) bool {
			return len(mounts[i].Root) > len(mounts[j].Root)
		})
	}

	return mounts[0], nil
}

// ParentPathForRootMount takes a container (child) and composes a path to
// the root mount point in the context of the parent.
func ParentPathForRootMount(parent IsolationResult, child IsolationResult) (*safepath.Path, error) {
	childRootMountInfo, err := MountInfoRoot(child)
	if err != nil {
		return nil, err
	}
	parentMountInfo, err := parentMountInfoFor(parent, childRootMountInfo)
	if err != nil {
		return nil, err
	}
	parentMountRoot, err := parent.MountRoot()
	if err != nil {
		return nil, err
	}
	path := parentMountRoot
	path, err = path.AppendAndResolveWithRelativeRoot(parentMountInfo.Mountpoint)
	if err != nil {
		return nil, err
	}
	return path.AppendAndResolveWithRelativeRoot(strings.TrimPrefix(childRootMountInfo.Root, parentMountInfo.Root))
}

func SafeJoin(res IsolationResult, elems ...string) (*safepath.Path, error) {
	mountRoot, err := res.MountRoot()
	if err != nil {
		return nil, err
	}
	return mountRoot.AppendAndResolveWithRelativeRoot(elems...)
}
<|endoftext|>


<commit_before>package lfs

import (
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

func ClearTempObjects() {
	filepath.Walk(LocalObjectTempDir, func(path string, info os.FileInfo, err error) error {
		if shouldDeleteTempObject(path, info) {
			return os.RemoveAll(path)
		}

		return err
	})
}

func shouldDeleteTempObject(path string, info os.FileInfo) bool {
	if info == nil || info.IsDir() {
		return false
	}

	base := filepath.Base(path)
	parts := strings.SplitN(base, "-", 2)
	oid := parts[0]
	if len(parts) < 2 || len(oid) != 64 {
		tracerx.Printf("Removing invalid tmp object file: %s", path)
		return true
	}

	if FileExists(localMediaPathNoCreate(oid)) {
		tracerx.Printf("Removing existing tmp object file: %s", path)
		return true
	}

	if time.Since(info.ModTime()) > time.Hour {
		tracerx.Printf("Removing old tmp object file: %s", path)
		return true
	}

	return false
}
<commit_msg>ララララー ラララー ララ<commit_after>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

func ClearTempObjects() {
	d, err := os.Open(LocalObjectTempDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening %q to clear old temp files: %s", LocalObjectTempDir, err)
		return
	}

	filenames, _ := d.Readdirnames(-1)
	for _, filename := range filenames {
		path := filepath.Join(LocalObjectTempDir, filename)
		if shouldDeleteTempObject(path) {
			os.RemoveAll(path)
		}
	}
}

func shouldDeleteTempObject(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}

	if info.IsDir() {
		return false
	}

	base := filepath.Base(path)
	parts := strings.SplitN(base, "-", 2)
	oid := parts[0]
	if len(parts) < 2 || len(oid) != 64 {
		tracerx.Printf("Removing invalid tmp object file: %s", path)
		return true
	}

	if FileExists(localMediaPathNoCreate(oid)) {
		tracerx.Printf("Removing existing tmp object file: %s", path)
		return true
	}

	if time.Since(info.ModTime()) > time.Hour {
		tracerx.Printf("Removing old tmp object file: %s", path)
		return true
	}

	return false
}
<|endoftext|>


<commit_before>//go:build !windows
// +build !windows

package net

import (
	"fmt"
	"net"
	gonetURL "net/url"
	"os"
	"strings"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	"github.com/containerd/cgroups"
	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/coreos/go-iptables/iptables"
)

const (
	/* "natsIsolationClassID" This is the integer value of the argument "0xb0540002", which is
	   b054:0002 . The major number (the left-hand side) is "BOSH", leet-ified.
	   The minor number (the right-hand side) is 2, indicating that this is the
	   second thing in our "BOSH" classid namespace.

	   _Hopefully_ noone uses a major number of "b054", and we avoid collisions _forever_!
	   If you need to select new classids for firewall rules or traffic control rules, keep
	   the major number "b054" for bosh stuff, unless there's a good reason to not.

	   The net_cls.classid structure is described in more detail here:
	   https://www.kernel.org/doc/Documentation/cgroup-v1/net_cls.txt
	*/
	natsIsolationClassID uint32 = 2958295042
)

// SetupNatsFirewall will setup the outgoing cgroup based rule that prevents everything except the agent to open connections to the nats api
func SetupNatsFirewall(mbus string) error {
	// return early if
	// we get a https url for mbus. case for create-env
	// we get an empty string. case for http_metadata_service (responsible to extract the agent-settings.json from the metadata endpoint)s
	if mbus == "" || strings.HasPrefix(mbus, "https://") {
		return nil
	}
	mbusURL, err := gonetURL.Parse(mbus)
	if err != nil || mbusURL.Hostname() == "" {
		return bosherr.WrapError(err, "Error parsing MbusURL")
	}

	host, port, err := net.SplitHostPort(mbusURL.Host)
	if err != nil {
		return bosherr.WrapError(err, "Error getting Port")
	}
	ipt, err := iptables.New()
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error")
	}
	exists, err := ipt.Exists("mangle", "POSTROUTING",
		"-d", "127.0.0.1",
		"-p", "tcp",
		"--dport", "2822",
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error checking for monit rule")
	}
	if !exists {
		err = ipt.Insert("mangle", "POSTROUTING", 1,
			"-d", "127.0.0.1",
			"-p", "tcp",
			"--dport", "2822",
			"-m", "cgroup",
			"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
			"-j", "ACCEPT",
		)
		if err != nil {
			return bosherr.WrapError(err, "Iptables Error inserting for monit rule")
		}
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for agent ACCEPT rule")
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-j", "DROP",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for non-agent DROP rule")
	}
	var isolationClassID = natsIsolationClassID
	natsAPICgroup, err := cgroups.New(cgroups.SingleSubsystem(cgroups.V1, cgroups.NetCLS), cgroups.StaticPath("/nats-api-access"), &specs.LinuxResources{
		Network: &specs.LinuxNetwork{
			ClassID: &isolationClassID,
		},
	})
	if err != nil {
		return bosherr.WrapError(err, "Error setting up cgroups for nats api access")
	}

	err = natsAPICgroup.AddProc(uint64(os.Getpid()), cgroups.NetCLS)
	return err
}
<commit_msg>Return early if v1 cgroups are not mounted<commit_after>//go:build !windows
// +build !windows

package net

import (
	"errors"
	"fmt"
	"net"
	gonetURL "net/url"
	"os"
	"strings"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	"github.com/containerd/cgroups"
	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/coreos/go-iptables/iptables"
)

const (
	/* "natsIsolationClassID" This is the integer value of the argument "0xb0540002", which is
	   b054:0002 . The major number (the left-hand side) is "BOSH", leet-ified.
	   The minor number (the right-hand side) is 2, indicating that this is the
	   second thing in our "BOSH" classid namespace.

	   _Hopefully_ noone uses a major number of "b054", and we avoid collisions _forever_!
	   If you need to select new classids for firewall rules or traffic control rules, keep
	   the major number "b054" for bosh stuff, unless there's a good reason to not.

	   The net_cls.classid structure is described in more detail here:
	   https://www.kernel.org/doc/Documentation/cgroup-v1/net_cls.txt
	*/
	natsIsolationClassID uint32 = 2958295042
)

// SetupNatsFirewall will setup the outgoing cgroup based rule that prevents everything except the agent to open connections to the nats api
func SetupNatsFirewall(mbus string) error {
	// return early if
	// we get a https url for mbus. case for create-env
	// we get an empty string. case for http_metadata_service (responsible to extract the agent-settings.json from the metadata endpoint)
	// we find that v1cgroups are not mounted (warden stemcells)
	if mbus == "" || strings.HasPrefix(mbus, "https://") {
		return nil
	}
	_, err := cgroups.V1()
	if err != nil {
		if errors.Is(err, cgroups.ErrMountPointNotExist) {
			return nil // v1cgroups are not mounted (warden stemcells)
		} else {
			return bosherr.WrapError(err, "Error retrieving cgroups mount point")
		}
	}
	mbusURL, err := gonetURL.Parse(mbus)
	if err != nil || mbusURL.Hostname() == "" {
		return bosherr.WrapError(err, "Error parsing MbusURL")
	}

	host, port, err := net.SplitHostPort(mbusURL.Host)
	if err != nil {
		return bosherr.WrapError(err, "Error getting Port")
	}
	ipt, err := iptables.New()
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error")
	}
	exists, err := ipt.Exists("mangle", "POSTROUTING",
		"-d", "127.0.0.1",
		"-p", "tcp",
		"--dport", "2822",
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error checking for monit rule")
	}
	if !exists {
		err = ipt.Insert("mangle", "POSTROUTING", 1,
			"-d", "127.0.0.1",
			"-p", "tcp",
			"--dport", "2822",
			"-m", "cgroup",
			"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
			"-j", "ACCEPT",
		)
		if err != nil {
			return bosherr.WrapError(err, "Iptables Error inserting for monit rule")
		}
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for agent ACCEPT rule")
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-j", "DROP",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for non-agent DROP rule")
	}
	var isolationClassID = natsIsolationClassID
	natsAPICgroup, err := cgroups.New(cgroups.SingleSubsystem(cgroups.V1, cgroups.NetCLS), cgroups.StaticPath("/nats-api-access"), &specs.LinuxResources{
		Network: &specs.LinuxNetwork{
			ClassID: &isolationClassID,
		},
	})
	if err != nil {
		return bosherr.WrapError(err, "Error setting up cgroups for nats api access")
	}

	err = natsAPICgroup.AddProc(uint64(os.Getpid()), cgroups.NetCLS)
	return err
}
<|endoftext|>


<commit_before>package librariesio

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

const (
	libraryVersion = "1"
	baseURL        = "https://libraries.io/api/"
	userAgent      = "go-librariesio/" + libraryVersion
	contentType    = "application/json"
	mediaType      = "application/json"
	timeout        = time.Second * 10
)

// Client for communicating with the libraries.io API
type Client struct {
	apiKey    string
	transport *http.Transport
	client    *http.Client
	UserAgent string
	BaseURL   *url.URL
}

// NewClient returns a new libraries.io API client
func NewClient(apiKey string) *Client {
	APIBaseURL, _ := url.Parse(baseURL)

	transport := &http.Transport{}
	client := &http.Client{Transport: transport}

	return &Client{
		apiKey:    apiKey,
		client:    client,
		transport: transport,
		UserAgent: userAgent,
		BaseURL:   APIBaseURL,
	}
}

// NewRequest creates a new API request
func (c *Client) NewRequest(method, urlStr string, data interface{}) (*http.Request, error) {
	relativeURL, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	absoluteURL := c.BaseURL.ResolveReference(relativeURL)

	var body io.ReadWriter
	if data != nil {
		body = new(bytes.Buffer)

		err := json.NewEncoder(body).Encode(data)
		if err != nil {
			return nil, err
		}
	}

	request, err := http.NewRequest(method, absoluteURL.String(), body)
	if err != nil {
		return nil, err
	}

	// set api_key for auth
	q := request.URL.Query()
	q.Set("api_key", c.apiKey)
	request.URL.RawQuery = q.Encode()

	request.Header.Set("Accept", mediaType)
	request.Header.Set("User-Agent", c.UserAgent)

	if data != nil {
		request.Header.Set("Content-Type", "application/json")
	}

	return request, nil
}

// ErrorResponse holds information about an unsuccesful API request
type ErrorResponse struct {
	Response *http.Response
	Message  string `json:"error"`
}

// Error interface implementation for ErrorResponse
func (r *ErrorResponse) Error() string {
	// Make sure to not show api_key
	q := r.Response.Request.URL.Query()
	q.Set("api_key", "REDACTED")
	r.Response.Request.URL.RawQuery = q.Encode()

	return fmt.Sprintf(
		"%v %v: %d %q",
		r.Response.Request.Method,
		r.Response.Request.URL,
		r.Response.StatusCode,
		r.Message,
	)
}

// CheckResponse checks the API response for errors and returns a ErrorResponse
// Responses are considered unsuccesful for status code other than 2xx.
func CheckResponse(resp *http.Response) error {
	if code := resp.StatusCode; 200 <= code && code <= 299 {
		return nil
	}

	errorResponse := &ErrorResponse{Response: resp}

	data, err := ioutil.ReadAll(resp.Body)
	if err == nil && data != nil {
		json.Unmarshal(data, errorResponse)
	}
	return errorResponse
}

// Do sends an HTTP request and returns an HTTP response
func (c *Client) Do(ctx context.Context, req *http.Request, v interface{}) (*http.Response, error) {
	req = req.WithContext(ctx)

	response, err := c.client.Do(req)
	if err != nil {
		select {
		case <-ctx.Done():
			c.transport.CancelRequest(req)
			return nil, ctx.Err()
		default:
			return nil, err
		}
	}
	defer response.Body.Close()

	err = CheckResponse(response)
	if err != nil {
		return response, err
	}

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

	err = json.Unmarshal(body, v)
	if err != nil {
		return nil, err
	}
	return response, nil
}
<commit_msg>Create redactAPIKey helper function<commit_after>package librariesio

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

const (
	libraryVersion = "1"
	baseURL        = "https://libraries.io/api/"
	userAgent      = "go-librariesio/" + libraryVersion
	contentType    = "application/json"
	mediaType      = "application/json"
	timeout        = time.Second * 10
)

// Client for communicating with the libraries.io API
type Client struct {
	apiKey    string
	transport *http.Transport
	client    *http.Client
	UserAgent string
	BaseURL   *url.URL
}

// NewClient returns a new libraries.io API client
func NewClient(apiKey string) *Client {
	APIBaseURL, _ := url.Parse(baseURL)

	transport := &http.Transport{}
	client := &http.Client{Transport: transport}

	return &Client{
		apiKey:    apiKey,
		client:    client,
		transport: transport,
		UserAgent: userAgent,
		BaseURL:   APIBaseURL,
	}
}

// NewRequest creates a new API request
func (c *Client) NewRequest(method, urlStr string, data interface{}) (*http.Request, error) {
	relativeURL, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	absoluteURL := c.BaseURL.ResolveReference(relativeURL)

	var body io.ReadWriter
	if data != nil {
		body = new(bytes.Buffer)

		err := json.NewEncoder(body).Encode(data)
		if err != nil {
			return nil, err
		}
	}

	request, err := http.NewRequest(method, absoluteURL.String(), body)
	if err != nil {
		return nil, err
	}

	// set api_key for auth
	q := request.URL.Query()
	q.Set("api_key", c.apiKey)
	request.URL.RawQuery = q.Encode()

	request.Header.Set("Accept", mediaType)
	request.Header.Set("User-Agent", c.UserAgent)

	if data != nil {
		request.Header.Set("Content-Type", "application/json")
	}

	return request, nil
}

// redactAPIKey overwrites the secret api_key query param
func redactAPIKey(url *url.URL) *url.URL {
	q := url.Query()
	q.Set("api_key", "REDACTED")
	url.RawQuery = q.Encode()
	return url
}

// ErrorResponse holds information about an unsuccesful API request
type ErrorResponse struct {
	Response *http.Response
	Message  string `json:"error"`
}

// Error interface implementation for ErrorResponse
func (r *ErrorResponse) Error() string {
	return fmt.Sprintf(
		"%v %v: %d %q",
		r.Response.Request.Method,
		redactAPIKey(r.Response.Request.URL),
		r.Response.StatusCode,
		r.Message,
	)
}

// CheckResponse checks the API response for errors and returns a ErrorResponse
// Responses are considered unsuccesful for status code other than 2xx.
func CheckResponse(resp *http.Response) error {
	if code := resp.StatusCode; 200 <= code && code <= 299 {
		return nil
	}

	errorResponse := &ErrorResponse{Response: resp}

	data, err := ioutil.ReadAll(resp.Body)
	if err == nil && data != nil {
		json.Unmarshal(data, errorResponse)
	}
	return errorResponse
}

// Do sends an HTTP request and returns an HTTP response
func (c *Client) Do(ctx context.Context, req *http.Request, v interface{}) (*http.Response, error) {
	req = req.WithContext(ctx)

	response, err := c.client.Do(req)
	if err != nil {
		select {
		case <-ctx.Done():
			c.transport.CancelRequest(req)
			return nil, ctx.Err()
		default:
			return nil, err
		}
	}
	defer response.Body.Close()

	err = CheckResponse(response)
	if err != nil {
		return response, err
	}

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

	err = json.Unmarshal(body, v)
	if err != nil {
		return nil, err
	}
	return response, nil
}
<|endoftext|>


<commit_before>package pgmgr

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/lib/pq"
)

// Migration stores a single migration's version and filename.
type Migration struct {
	Filename string
	Version  int64
}

// Migration directions
const (
	DOWN = iota
	UP
)

const datetimeFormat = "20060102130405"

// Create creates the database specified by the configuration.
func Create(c *Config) error {
	return sh("createdb", []string{c.Database})
}

// Drop drops the database specified by the configuration.
func Drop(c *Config) error {
	return sh("dropdb", []string{c.Database})
}

// Dump dumps the schema and contents of the database to the dump file.
func Dump(c *Config) error {
	// dump schema first...
	schema, err := shRead("pg_dump", []string{"--schema-only", c.Database})
	if err != nil {
		return err
	}

	// then selected data...
	args := []string{c.Database, "--data-only"}
	if len(c.SeedTables) > 0 {
		for _, table := range c.SeedTables {
			println("pulling data for", table)
			args = append(args, "-t", table)
		}
	}
	println(strings.Join(args, ""))

	seeds, err := shRead("pg_dump", args)
	if err != nil {
		return err
	}

	// and combine into one file.
	file, err := os.OpenFile(c.DumpFile, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0600)
	if err != nil {
		return err
	}

	file.Write(*schema)
	file.Write(*seeds)
	file.Close()

	return nil
}

// Load loads the database from the dump file using psql.
func Load(c *Config) error {
	return sh("psql", []string{"-d", c.Database, "-f", c.DumpFile})
}

// Migrate applies un-applied migrations in the specified MigrationFolder.
func Migrate(c *Config) error {
	migrations, err := migrations(c, "up")
	if err != nil {
		return err
	}

	// ensure the version table is created
	_, err = getOrInitializeVersion(c)
	if err != nil {
		return err
	}

	appliedAny := false
	for _, m := range migrations {
		if applied, _ := migrationIsApplied(c, m.Version); !applied {
			fmt.Println("== Applying", m.Filename, "==")
			t0 := time.Now()

			if err = applyMigration(c, m, UP); err != nil { // halt the migration process and return the error.
				fmt.Println(err)
				fmt.Println("")
				fmt.Println("ERROR! Aborting the migration process.")
				return err
			}

			fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds()/1e6, "ms ==")
			appliedAny = true
		}
	}

	if !appliedAny {
		fmt.Println("Nothing to do; all migrations already applied.")
	}

	return nil
}

// Rollback un-applies the latest migration, if possible.
func Rollback(c *Config) error {
	migrations, err := migrations(c, "down")
	if err != nil {
		return err
	}

	v, _ := Version(c)
	var toRollback *Migration
	for _, m := range migrations {
		if m.Version == v {
			toRollback = &m
			break
		}
	}

	if toRollback == nil {
		return nil
	}

	// rollback only the last migration
	fmt.Println("== Reverting", toRollback.Filename, "==")
	t0 := time.Now()

	if err = applyMigration(c, *toRollback, DOWN); err != nil {
		return err
	}

	fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds()/1e6, "ms ==")

	return nil
}

// Version returns the highest version number stored in the database. This is not
// necessarily enough info to uniquely identify the version, since there may
// be backdated migrations which have not yet applied.
func Version(c *Config) (int64, error) {
	db, err := openConnection(c)
	defer db.Close()

	if err != nil {
		return -1, err
	}

	// if the table doesn't exist, we're at version -1
	var hasTable bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM pg_catalog.pg_tables WHERE tablename='schema_migrations')").Scan(&hasTable)
	if hasTable != true {
		return -1, err
	}

	var version int64
	err = db.QueryRow("SELECT COALESCE(MAX(version)::text, '-1') FROM schema_migrations").Scan(&version)

	return version, err
}

// Initialize creates the schema_migrations table on what should be a new database.
func Initialize(c *Config) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	var statement string
	if c.ColumnType == "string" {
		statement = "CREATE TABLE schema_migrations (version CHARACTER VARYING(255) NOT NULL)"
	} else {
		statement = "CREATE TABLE schema_migrations (version INTEGER NOT NULL)"
	}

	_, err = db.Exec(statement)
	if err != nil {
		return err
	}

	return nil
}

// CreateMigration generates new, empty migration files.
func CreateMigration(c *Config, name string) error {
	version := generateVersion(c)
	upFilepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".up.sql"))
	downFilepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".down.sql"))

	err := ioutil.WriteFile(upFilepath, []byte(`-- Migration goes here.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", upFilepath)

	err = ioutil.WriteFile(downFilepath, []byte(`-- Rollback of migration goes here. If you don't want to write it, delete this file.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", downFilepath)

	return nil
}

func generateVersion(c *Config) string {
	// TODO: guarantee no conflicts by incrementing if there is a conflict
	t := time.Now()

	if c.Format == "datetime" {
		return t.Format(datetimeFormat)
	}

	return strconv.FormatInt(t.Unix(), 10)
}

// need access to the original query contents in order to print it out properly,
// unfortunately.
func formatPgErr(contents *[]byte, pgerr *pq.Error) string {
	pos, _ := strconv.Atoi(pgerr.Position)
	lineNo := bytes.Count((*contents)[:pos], []byte("\n")) + 1
	columnNo := pos - bytes.LastIndex((*contents)[:pos], []byte("\n")) - 1

	return fmt.Sprint("PGERROR: line ", lineNo, " pos ", columnNo, ": ", pgerr.Message, ". ", pgerr.Detail)
}

func applyMigration(c *Config, m Migration, direction int) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	contents, err := ioutil.ReadFile(filepath.Join(c.MigrationFolder, m.Filename))
	if err != nil {
		return err
	}

	if _, err = tx.Exec(string(contents)); err != nil {
		tx.Rollback()
		return errors.New(formatPgErr(&contents, err.(*pq.Error)))
	}

	if direction == UP {
		if err = insertSchemaVersion(c, tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	} else {
		if err = deleteSchemaVersion(c, tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	}

	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

func insertSchemaVersion(c *Config, tx *sql.Tx, version int64) error {
	_, err := tx.Exec("INSERT INTO schema_migrations (version) VALUES ($1) RETURNING version", typedVersion(c, version))
	return err
}

func deleteSchemaVersion(c *Config, tx *sql.Tx, version int64) error {
	_, err := tx.Exec("DELETE FROM schema_migrations WHERE version = $1", typedVersion(c, version))
	return err
}

func typedVersion(c *Config, version int64) interface{} {
	if c.ColumnType == "string" {
		return strconv.FormatInt(version, 10)
	}
	return version
}

func migrationIsApplied(c *Config, version int64) (bool, error) {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return false, err
	}

	var applied bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM schema_migrations WHERE version = $1);", version).Scan(&applied)
	if err != nil {
		return false, err
	}

	return applied, nil
}

func getOrInitializeVersion(c *Config) (int64, error) {
	var v int64
	if v, _ = Version(c); v < 0 {
		if err := Initialize(c); err != nil {
			return -1, err
		}
	}

	return v, nil
}

func openConnection(c *Config) (*sql.DB, error) {
	db, err := sql.Open("postgres", sqlConnectionString(c))
	return db, err
}

func sqlConnectionString(c *Config) string {
	return fmt.Sprint(
		" user='", c.Username, "'",
		" dbname='", c.Database, "'",
		" password='", c.Password, "'",
		" host='", c.Host, "'",
		" port=", c.Port,
		" sslmode=", "disable")
}

func migrations(c *Config, direction string) ([]Migration, error) {
	re := regexp.MustCompile("^[0-9]+")

	migrations := []Migration{}
	files, err := ioutil.ReadDir(c.MigrationFolder)
	if err != nil {
		return migrations, err
	}

	for _, file := range files {
		if match, _ := regexp.MatchString("[0-9]+_.+."+direction+".sql", file.Name()); match {
			version, _ := strconv.ParseInt(re.FindString(file.Name()), 10, 64)
			migrations = append(migrations, Migration{Filename: file.Name(), Version: version})
		}
	}

	return migrations, nil
}

func sh(command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	fmt.Println(string(output))
	if err != nil {
		return err
	}

	return nil
}

func shRead(command string, args []string) (*[]byte, error) {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	return &output, err
}
<commit_msg>Add unique constraint on schema_migrations.version<commit_after>package pgmgr

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/lib/pq"
)

// Migration stores a single migration's version and filename.
type Migration struct {
	Filename string
	Version  int64
}

// Migration directions
const (
	DOWN = iota
	UP
)

const datetimeFormat = "20060102130405"

// Create creates the database specified by the configuration.
func Create(c *Config) error {
	return sh("createdb", []string{c.Database})
}

// Drop drops the database specified by the configuration.
func Drop(c *Config) error {
	return sh("dropdb", []string{c.Database})
}

// Dump dumps the schema and contents of the database to the dump file.
func Dump(c *Config) error {
	// dump schema first...
	schema, err := shRead("pg_dump", []string{"--schema-only", c.Database})
	if err != nil {
		return err
	}

	// then selected data...
	args := []string{c.Database, "--data-only"}
	if len(c.SeedTables) > 0 {
		for _, table := range c.SeedTables {
			println("pulling data for", table)
			args = append(args, "-t", table)
		}
	}
	println(strings.Join(args, ""))

	seeds, err := shRead("pg_dump", args)
	if err != nil {
		return err
	}

	// and combine into one file.
	file, err := os.OpenFile(c.DumpFile, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0600)
	if err != nil {
		return err
	}

	file.Write(*schema)
	file.Write(*seeds)
	file.Close()

	return nil
}

// Load loads the database from the dump file using psql.
func Load(c *Config) error {
	return sh("psql", []string{"-d", c.Database, "-f", c.DumpFile})
}

// Migrate applies un-applied migrations in the specified MigrationFolder.
func Migrate(c *Config) error {
	migrations, err := migrations(c, "up")
	if err != nil {
		return err
	}

	// ensure the version table is created
	_, err = getOrInitializeVersion(c)
	if err != nil {
		return err
	}

	appliedAny := false
	for _, m := range migrations {
		if applied, _ := migrationIsApplied(c, m.Version); !applied {
			fmt.Println("== Applying", m.Filename, "==")
			t0 := time.Now()

			if err = applyMigration(c, m, UP); err != nil { // halt the migration process and return the error.
				fmt.Println(err)
				fmt.Println("")
				fmt.Println("ERROR! Aborting the migration process.")
				return err
			}

			fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds()/1e6, "ms ==")
			appliedAny = true
		}
	}

	if !appliedAny {
		fmt.Println("Nothing to do; all migrations already applied.")
	}

	return nil
}

// Rollback un-applies the latest migration, if possible.
func Rollback(c *Config) error {
	migrations, err := migrations(c, "down")
	if err != nil {
		return err
	}

	v, _ := Version(c)
	var toRollback *Migration
	for _, m := range migrations {
		if m.Version == v {
			toRollback = &m
			break
		}
	}

	if toRollback == nil {
		return nil
	}

	// rollback only the last migration
	fmt.Println("== Reverting", toRollback.Filename, "==")
	t0 := time.Now()

	if err = applyMigration(c, *toRollback, DOWN); err != nil {
		return err
	}

	fmt.Println("== Completed in", time.Now().Sub(t0).Nanoseconds()/1e6, "ms ==")

	return nil
}

// Version returns the highest version number stored in the database. This is not
// necessarily enough info to uniquely identify the version, since there may
// be backdated migrations which have not yet applied.
func Version(c *Config) (int64, error) {
	db, err := openConnection(c)
	defer db.Close()

	if err != nil {
		return -1, err
	}

	// if the table doesn't exist, we're at version -1
	var hasTable bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM pg_catalog.pg_tables WHERE tablename='schema_migrations')").Scan(&hasTable)
	if hasTable != true {
		return -1, err
	}

	var version int64
	err = db.QueryRow("SELECT COALESCE(MAX(version)::text, '-1') FROM schema_migrations").Scan(&version)

	return version, err
}

// Initialize creates the schema_migrations table on what should be a new database.
func Initialize(c *Config) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	var statement string
	if c.ColumnType == "string" {
		statement = "CREATE TABLE schema_migrations (version CHARACTER VARYING(255) NOT NULL UNIQUE)"
	} else {
		statement = "CREATE TABLE schema_migrations (version INTEGER NOT NULL UNIQUE)"
	}

	_, err = db.Exec(statement)
	if err != nil {
		return err
	}

	return nil
}

// CreateMigration generates new, empty migration files.
func CreateMigration(c *Config, name string) error {
	version := generateVersion(c)
	upFilepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".up.sql"))
	downFilepath := filepath.Join(c.MigrationFolder, fmt.Sprint(version, "_", name, ".down.sql"))

	err := ioutil.WriteFile(upFilepath, []byte(`-- Migration goes here.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", upFilepath)

	err = ioutil.WriteFile(downFilepath, []byte(`-- Rollback of migration goes here. If you don't want to write it, delete this file.`), 0644)
	if err != nil {
		return err
	}
	fmt.Println("Created", downFilepath)

	return nil
}

func generateVersion(c *Config) string {
	// TODO: guarantee no conflicts by incrementing if there is a conflict
	t := time.Now()

	if c.Format == "datetime" {
		return t.Format(datetimeFormat)
	}

	return strconv.FormatInt(t.Unix(), 10)
}

// need access to the original query contents in order to print it out properly,
// unfortunately.
func formatPgErr(contents *[]byte, pgerr *pq.Error) string {
	pos, _ := strconv.Atoi(pgerr.Position)
	lineNo := bytes.Count((*contents)[:pos], []byte("\n")) + 1
	columnNo := pos - bytes.LastIndex((*contents)[:pos], []byte("\n")) - 1

	return fmt.Sprint("PGERROR: line ", lineNo, " pos ", columnNo, ": ", pgerr.Message, ". ", pgerr.Detail)
}

func applyMigration(c *Config, m Migration, direction int) error {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	contents, err := ioutil.ReadFile(filepath.Join(c.MigrationFolder, m.Filename))
	if err != nil {
		return err
	}

	if _, err = tx.Exec(string(contents)); err != nil {
		tx.Rollback()
		return errors.New(formatPgErr(&contents, err.(*pq.Error)))
	}

	if direction == UP {
		if err = insertSchemaVersion(c, tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	} else {
		if err = deleteSchemaVersion(c, tx, m.Version); err != nil {
			tx.Rollback()
			return errors.New(formatPgErr(&contents, err.(*pq.Error)))
		}
	}

	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

func insertSchemaVersion(c *Config, tx *sql.Tx, version int64) error {
	_, err := tx.Exec("INSERT INTO schema_migrations (version) VALUES ($1) RETURNING version", typedVersion(c, version))
	return err
}

func deleteSchemaVersion(c *Config, tx *sql.Tx, version int64) error {
	_, err := tx.Exec("DELETE FROM schema_migrations WHERE version = $1", typedVersion(c, version))
	return err
}

func typedVersion(c *Config, version int64) interface{} {
	if c.ColumnType == "string" {
		return strconv.FormatInt(version, 10)
	}
	return version
}

func migrationIsApplied(c *Config, version int64) (bool, error) {
	db, err := openConnection(c)
	defer db.Close()
	if err != nil {
		return false, err
	}

	var applied bool
	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM schema_migrations WHERE version = $1);", version).Scan(&applied)
	if err != nil {
		return false, err
	}

	return applied, nil
}

func getOrInitializeVersion(c *Config) (int64, error) {
	var v int64
	if v, _ = Version(c); v < 0 {
		if err := Initialize(c); err != nil {
			return -1, err
		}
	}

	return v, nil
}

func openConnection(c *Config) (*sql.DB, error) {
	db, err := sql.Open("postgres", sqlConnectionString(c))
	return db, err
}

func sqlConnectionString(c *Config) string {
	return fmt.Sprint(
		" user='", c.Username, "'",
		" dbname='", c.Database, "'",
		" password='", c.Password, "'",
		" host='", c.Host, "'",
		" port=", c.Port,
		" sslmode=", "disable")
}

func migrations(c *Config, direction string) ([]Migration, error) {
	re := regexp.MustCompile("^[0-9]+")

	migrations := []Migration{}
	files, err := ioutil.ReadDir(c.MigrationFolder)
	if err != nil {
		return migrations, err
	}

	for _, file := range files {
		if match, _ := regexp.MatchString("[0-9]+_.+."+direction+".sql", file.Name()); match {
			version, _ := strconv.ParseInt(re.FindString(file.Name()), 10, 64)
			migrations = append(migrations, Migration{Filename: file.Name(), Version: version})
		}
	}

	return migrations, nil
}

func sh(command string, args []string) error {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	fmt.Println(string(output))
	if err != nil {
		return err
	}

	return nil
}

func shRead(command string, args []string) (*[]byte, error) {
	c := exec.Command(command, args...)
	output, err := c.CombinedOutput()
	return &output, err
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 common

import (
	"net/http"
	"strings"

	"github.com/emicklei/go-restful"
	"github.com/go-openapi/spec"
)

// OpenAPIDefinition describes single type. Normally these definitions are auto-generated using gen-openapi.
type OpenAPIDefinition struct {
	Schema       spec.Schema
	Dependencies []string
}

type ReferenceCallback func(path string) spec.Ref

// OpenAPIDefinitions is collection of all definitions.
type GetOpenAPIDefinitions func(ReferenceCallback) map[string]OpenAPIDefinition

// OpenAPIDefinitionGetter gets openAPI definitions for a given type. If a type implements this interface,
// the definition returned by it will be used, otherwise the auto-generated definitions will be used. See
// GetOpenAPITypeFormat for more information about trade-offs of using this interface or GetOpenAPITypeFormat method when
// possible.
type OpenAPIDefinitionGetter interface {
	OpenAPIDefinition() *OpenAPIDefinition
}

type PathHandler interface {
	Handle(path string, handler http.Handler)
}

// Config is set of configuration for openAPI spec generation.
type Config struct {
	// List of supported protocols such as https, http, etc.
	ProtocolList []string

	// Info is general information about the API.
	Info *spec.Info

	// DefaultResponse will be used if an operation does not have any responses listed. It
	// will show up as ... "responses" : {"default" : $DefaultResponse} in the spec.
	DefaultResponse *spec.Response

	// CommonResponses will be added as a response to all operation specs. This is a good place to add common
	// responses such as authorization failed.
	CommonResponses map[int]spec.Response

	// List of webservice's path prefixes to ignore
	IgnorePrefixes []string

	// OpenAPIDefinitions should provide definition for all models used by routes. Failure to provide this map
	// or any of the models will result in spec generation failure.
	GetDefinitions GetOpenAPIDefinitions

	// GetOperationIDAndTags returns operation id and tags for a restful route. It is an optional function to customize operation IDs.
	GetOperationIDAndTags func(r *restful.Route) (string, []string, error)

	// GetDefinitionName returns a friendly name for a definition base on the serving path. parameter `name` is the full name of the definition.
	// It is an optional function to customize model names.
	GetDefinitionName func(name string) (string, spec.Extensions)

	// PostProcessSpec runs after the spec is ready to serve. It allows a final modification to the spec before serving.
	PostProcessSpec func(*spec.Swagger) (*spec.Swagger, error)

	// SecurityDefinitions is list of all security definitions for OpenAPI service. If this is not nil, the user of config
	// is responsible to provide DefaultSecurity and (maybe) add unauthorized response to CommonResponses.
	SecurityDefinitions *spec.SecurityDefinitions

	// DefaultSecurity for all operations. This will pass as spec.SwaggerProps.Security to OpenAPI.
	// For most cases, this will be list of acceptable definitions in SecurityDefinitions.
	DefaultSecurity []map[string][]string
}

// This function is a reference for converting go (or any custom type) to a simple open API type,format pair. There are
// two ways to customize spec for a type. If you add it here, a type will be converted to a simple type and the type
// comment (the comment that is added before type definition) will be lost. The spec will still have the property
// comment. The second way is to implement OpenAPIDefinitionGetter interface. That function can customize the spec (so
// the spec does not need to be simple type,format) or can even return a simple type,format (e.g. IntOrString). For simple
// type formats, the benefit of adding OpenAPIDefinitionGetter interface is to keep both type and property documentation.
// Example:
// type Sample struct {
//      ...
//      // port of the server
//      port IntOrString
//      ...
// }
// // IntOrString documentation...
// type IntOrString { ... }
//
// Adding IntOrString to this function:
// "port" : {
//           format:      "string",
//           type:        "int-or-string",
//           Description: "port of the server"
// }
//
// Implement OpenAPIDefinitionGetter for IntOrString:
//
// "port" : {
//           $Ref:    "#/definitions/IntOrString"
//           Description: "port of the server"
// }
// ...
// definitions:
// {
//           "IntOrString": {
//                     format:      "string",
//                     type:        "int-or-string",
//                     Description: "IntOrString documentation..."    // new
//           }
// }
//
func GetOpenAPITypeFormat(typeName string) (string, string) {
	schemaTypeFormatMap := map[string][]string{
		"uint":        {"integer", "int32"},
		"uint8":       {"integer", "byte"},
		"uint16":      {"integer", "int32"},
		"uint32":      {"integer", "int64"},
		"uint64":      {"integer", "int64"},
		"int":         {"integer", "int32"},
		"int8":        {"integer", "byte"},
		"int16":       {"integer", "int32"},
		"int32":       {"integer", "int32"},
		"int64":       {"integer", "int64"},
		"byte":        {"integer", "byte"},
		"float64":     {"number", "double"},
		"float32":     {"number", "float"},
		"bool":        {"boolean", ""},
		"time.Time":   {"string", "date-time"},
		"string":      {"string", ""},
		"integer":     {"integer", ""},
		"number":      {"number", ""},
		"boolean":     {"boolean", ""},
		"[]byte":      {"string", "byte"}, // base64 encoded characters
		"interface{}": {"object", ""},
	}
	mapped, ok := schemaTypeFormatMap[typeName]
	if !ok {
		return "", ""
	}
	return mapped[0], mapped[1]
}

func EscapeJsonPointer(p string) string {
	// Escaping reference name using rfc6901
	p = strings.Replace(p, "~", "~0", -1)
	p = strings.Replace(p, "/", "~1", -1)
	return p
}
<commit_msg>initialise schemaTypeFormatMap once at package level instead of every time GetOpenAPITypeFormat is called<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 common

import (
	"net/http"
	"strings"

	"github.com/emicklei/go-restful"
	"github.com/go-openapi/spec"
)

// OpenAPIDefinition describes single type. Normally these definitions are auto-generated using gen-openapi.
type OpenAPIDefinition struct {
	Schema       spec.Schema
	Dependencies []string
}

type ReferenceCallback func(path string) spec.Ref

// OpenAPIDefinitions is collection of all definitions.
type GetOpenAPIDefinitions func(ReferenceCallback) map[string]OpenAPIDefinition

// OpenAPIDefinitionGetter gets openAPI definitions for a given type. If a type implements this interface,
// the definition returned by it will be used, otherwise the auto-generated definitions will be used. See
// GetOpenAPITypeFormat for more information about trade-offs of using this interface or GetOpenAPITypeFormat method when
// possible.
type OpenAPIDefinitionGetter interface {
	OpenAPIDefinition() *OpenAPIDefinition
}

type PathHandler interface {
	Handle(path string, handler http.Handler)
}

// Config is set of configuration for openAPI spec generation.
type Config struct {
	// List of supported protocols such as https, http, etc.
	ProtocolList []string

	// Info is general information about the API.
	Info *spec.Info

	// DefaultResponse will be used if an operation does not have any responses listed. It
	// will show up as ... "responses" : {"default" : $DefaultResponse} in the spec.
	DefaultResponse *spec.Response

	// CommonResponses will be added as a response to all operation specs. This is a good place to add common
	// responses such as authorization failed.
	CommonResponses map[int]spec.Response

	// List of webservice's path prefixes to ignore
	IgnorePrefixes []string

	// OpenAPIDefinitions should provide definition for all models used by routes. Failure to provide this map
	// or any of the models will result in spec generation failure.
	GetDefinitions GetOpenAPIDefinitions

	// GetOperationIDAndTags returns operation id and tags for a restful route. It is an optional function to customize operation IDs.
	GetOperationIDAndTags func(r *restful.Route) (string, []string, error)

	// GetDefinitionName returns a friendly name for a definition base on the serving path. parameter `name` is the full name of the definition.
	// It is an optional function to customize model names.
	GetDefinitionName func(name string) (string, spec.Extensions)

	// PostProcessSpec runs after the spec is ready to serve. It allows a final modification to the spec before serving.
	PostProcessSpec func(*spec.Swagger) (*spec.Swagger, error)

	// SecurityDefinitions is list of all security definitions for OpenAPI service. If this is not nil, the user of config
	// is responsible to provide DefaultSecurity and (maybe) add unauthorized response to CommonResponses.
	SecurityDefinitions *spec.SecurityDefinitions

	// DefaultSecurity for all operations. This will pass as spec.SwaggerProps.Security to OpenAPI.
	// For most cases, this will be list of acceptable definitions in SecurityDefinitions.
	DefaultSecurity []map[string][]string
}

var schemaTypeFormatMap = map[string][]string{
	"uint":        {"integer", "int32"},
	"uint8":       {"integer", "byte"},
	"uint16":      {"integer", "int32"},
	"uint32":      {"integer", "int64"},
	"uint64":      {"integer", "int64"},
	"int":         {"integer", "int32"},
	"int8":        {"integer", "byte"},
	"int16":       {"integer", "int32"},
	"int32":       {"integer", "int32"},
	"int64":       {"integer", "int64"},
	"byte":        {"integer", "byte"},
	"float64":     {"number", "double"},
	"float32":     {"number", "float"},
	"bool":        {"boolean", ""},
	"time.Time":   {"string", "date-time"},
	"string":      {"string", ""},
	"integer":     {"integer", ""},
	"number":      {"number", ""},
	"boolean":     {"boolean", ""},
	"[]byte":      {"string", "byte"}, // base64 encoded characters
	"interface{}": {"object", ""},
}

// This function is a reference for converting go (or any custom type) to a simple open API type,format pair. There are
// two ways to customize spec for a type. If you add it here, a type will be converted to a simple type and the type
// comment (the comment that is added before type definition) will be lost. The spec will still have the property
// comment. The second way is to implement OpenAPIDefinitionGetter interface. That function can customize the spec (so
// the spec does not need to be simple type,format) or can even return a simple type,format (e.g. IntOrString). For simple
// type formats, the benefit of adding OpenAPIDefinitionGetter interface is to keep both type and property documentation.
// Example:
// type Sample struct {
//      ...
//      // port of the server
//      port IntOrString
//      ...
// }
// // IntOrString documentation...
// type IntOrString { ... }
//
// Adding IntOrString to this function:
// "port" : {
//           format:      "string",
//           type:        "int-or-string",
//           Description: "port of the server"
// }
//
// Implement OpenAPIDefinitionGetter for IntOrString:
//
// "port" : {
//           $Ref:    "#/definitions/IntOrString"
//           Description: "port of the server"
// }
// ...
// definitions:
// {
//           "IntOrString": {
//                     format:      "string",
//                     type:        "int-or-string",
//                     Description: "IntOrString documentation..."    // new
//           }
// }
//
func GetOpenAPITypeFormat(typeName string) (string, string) {
	mapped, ok := schemaTypeFormatMap[typeName]
	if !ok {
		return "", ""
	}
	return mapped[0], mapped[1]
}

func EscapeJsonPointer(p string) string {
	// Escaping reference name using rfc6901
	p = strings.Replace(p, "~", "~0", -1)
	p = strings.Replace(p, "/", "~1", -1)
	return p
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"flag"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"
	"time"
)

import "git.torproject.org/pluggable-transports/goptlib.git"

const ptMethodName = "meek"
const sessionIdLength = 32
const initPollInterval = 10 * time.Millisecond
const maxPollInterval = 5 * time.Second
const pollIntervalMultiplier = 1.5

var ptInfo pt.ClientInfo

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func roundTrip(u, sessionId string, buf []byte) (*http.Response, error) {
	req, err := http.NewRequest("POST", u, bytes.NewReader(buf))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Session-Id", sessionId)
	return http.DefaultClient.Do(req)
}

func copyLoop(conn net.Conn, u, sessionId string) error {
	buf := make([]byte, 0x10000)
	var interval time.Duration

	conn.SetReadDeadline(time.Now().Add(initPollInterval))
	for {
		nr, readErr := conn.Read(buf)
		// log.Printf("read from local: %q", buf[:nr])

		resp, err := roundTrip(u, sessionId, buf[:nr])
		if err != nil {
			return err
		}

		nw, err := io.Copy(conn, resp.Body)
		if err != nil {
			return err
		}
		// log.Printf("read from remote: %d", nw)

		if readErr != nil {
			if e, ok := readErr.(net.Error); !ok || !e.Timeout() {
				return readErr
			}
		}

		if nw > 0 {
			interval = 0
		} else if interval < initPollInterval {
			interval = initPollInterval
		} else {
			interval = time.Duration(float64(interval) * pollIntervalMultiplier)
		}
		if interval > maxPollInterval {
			interval = maxPollInterval
		}
		// log.Printf("next poll %.6f s", interval.Seconds())
		conn.SetReadDeadline(time.Now().Add(interval))
	}

	return nil
}

func genSessionId() string {
	buf := make([]byte, sessionIdLength)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err.Error())
	}
	return base64.StdEncoding.EncodeToString(buf)
}

func handler(conn *pt.SocksConn) error {
	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	defer conn.Close()
	err := conn.Grant(&net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: 0})
	if err != nil {
		return err
	}

	sessionId := genSessionId()
	u := url.URL{
		Scheme: "http",
		Host:   conn.Req.Target,
		Path:   "/",
	}

	return copyLoop(conn, u.String(), sessionId)
}

func acceptLoop(ln *pt.SocksListener) error {
	defer ln.Close()
	for {
		conn, err := ln.AcceptSocks()
		if err != nil {
			log.Printf("error in AcceptSocks: %s", err)
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				return err
			}
			continue
		}
		go func() {
			err := handler(conn)
			if err != nil {
				log.Printf("error in handling request: %s", err)
			}
		}()
	}
}

func main() {
	var logFilename string

	flag.StringVar(&logFilename, "log", "", "name of log file")
	flag.Parse()

	if logFilename != "" {
		f, err := os.OpenFile(logFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			log.Fatalf("error opening log file: %s", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	var err error
	ptInfo, err = pt.ClientSetup([]string{ptMethodName})
	if err != nil {
		log.Fatalf("error in ClientSetup: %s", err)
	}

	listeners := make([]net.Listener, 0)
	for _, methodName := range ptInfo.MethodNames {
		switch methodName {
		case ptMethodName:
			ln, err := pt.ListenSocks("tcp", "127.0.0.1:0")
			if err != nil {
				pt.CmethodError(methodName, err.Error())
				break
			}
			go acceptLoop(ln)
			pt.Cmethod(methodName, ln.Version(), ln.Addr())
			log.Printf("listening on %s", ln.Addr())
			listeners = append(listeners, ln)
		default:
			pt.CmethodError(methodName, "no such method")
		}
	}
	pt.CmethodsDone()

	var numHandlers int = 0
	var sig os.Signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// wait for first signal
	sig = nil
	for sig == nil {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}
	for _, ln := range listeners {
		ln.Close()
	}

	if sig == syscall.SIGTERM {
		return
	}

	// wait for second signal or no more handlers
	sig = nil
	for sig == nil && numHandlers != 0 {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}

	log.Printf("done")
}
<commit_msg>Increase initial poll interval to 100 ms from 10 ms.<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"flag"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"
	"time"
)

import "git.torproject.org/pluggable-transports/goptlib.git"

const ptMethodName = "meek"
const sessionIdLength = 32
const initPollInterval = 100 * time.Millisecond
const maxPollInterval = 5 * time.Second
const pollIntervalMultiplier = 1.5

var ptInfo pt.ClientInfo

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func roundTrip(u, sessionId string, buf []byte) (*http.Response, error) {
	req, err := http.NewRequest("POST", u, bytes.NewReader(buf))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Session-Id", sessionId)
	return http.DefaultClient.Do(req)
}

func copyLoop(conn net.Conn, u, sessionId string) error {
	buf := make([]byte, 0x10000)
	var interval time.Duration

	conn.SetReadDeadline(time.Now().Add(initPollInterval))
	for {
		nr, readErr := conn.Read(buf)
		// log.Printf("read from local: %q", buf[:nr])

		resp, err := roundTrip(u, sessionId, buf[:nr])
		if err != nil {
			return err
		}

		nw, err := io.Copy(conn, resp.Body)
		if err != nil {
			return err
		}
		// log.Printf("read from remote: %d", nw)

		if readErr != nil {
			if e, ok := readErr.(net.Error); !ok || !e.Timeout() {
				return readErr
			}
		}

		if nw > 0 {
			interval = 0
		} else if interval < initPollInterval {
			interval = initPollInterval
		} else {
			interval = time.Duration(float64(interval) * pollIntervalMultiplier)
		}
		if interval > maxPollInterval {
			interval = maxPollInterval
		}
		// log.Printf("next poll %.6f s", interval.Seconds())
		conn.SetReadDeadline(time.Now().Add(interval))
	}

	return nil
}

func genSessionId() string {
	buf := make([]byte, sessionIdLength)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err.Error())
	}
	return base64.StdEncoding.EncodeToString(buf)
}

func handler(conn *pt.SocksConn) error {
	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	defer conn.Close()
	err := conn.Grant(&net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: 0})
	if err != nil {
		return err
	}

	sessionId := genSessionId()
	u := url.URL{
		Scheme: "http",
		Host:   conn.Req.Target,
		Path:   "/",
	}

	return copyLoop(conn, u.String(), sessionId)
}

func acceptLoop(ln *pt.SocksListener) error {
	defer ln.Close()
	for {
		conn, err := ln.AcceptSocks()
		if err != nil {
			log.Printf("error in AcceptSocks: %s", err)
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				return err
			}
			continue
		}
		go func() {
			err := handler(conn)
			if err != nil {
				log.Printf("error in handling request: %s", err)
			}
		}()
	}
}

func main() {
	var logFilename string

	flag.StringVar(&logFilename, "log", "", "name of log file")
	flag.Parse()

	if logFilename != "" {
		f, err := os.OpenFile(logFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			log.Fatalf("error opening log file: %s", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	var err error
	ptInfo, err = pt.ClientSetup([]string{ptMethodName})
	if err != nil {
		log.Fatalf("error in ClientSetup: %s", err)
	}

	listeners := make([]net.Listener, 0)
	for _, methodName := range ptInfo.MethodNames {
		switch methodName {
		case ptMethodName:
			ln, err := pt.ListenSocks("tcp", "127.0.0.1:0")
			if err != nil {
				pt.CmethodError(methodName, err.Error())
				break
			}
			go acceptLoop(ln)
			pt.Cmethod(methodName, ln.Version(), ln.Addr())
			log.Printf("listening on %s", ln.Addr())
			listeners = append(listeners, ln)
		default:
			pt.CmethodError(methodName, "no such method")
		}
	}
	pt.CmethodsDone()

	var numHandlers int = 0
	var sig os.Signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// wait for first signal
	sig = nil
	for sig == nil {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}
	for _, ln := range listeners {
		ln.Close()
	}

	if sig == syscall.SIGTERM {
		return
	}

	// wait for second signal or no more handlers
	sig = nil
	for sig == nil && numHandlers != 0 {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}

	log.Printf("done")
}
<|endoftext|>