<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

var bars []*uiprogress.Bar

// runClient spawns threads for parallel uplink/downlink via TCP
func runClient(connectionType string, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"codePhrase": codePhrase,
		"connection": connectionType,
	})
	var wg sync.WaitGroup
	wg.Add(numberConnections)

	uiprogress.Start()
	if !debugFlag {
		bars = make([]*uiprogress.Bar, numberConnections)
	}
	for id := 0; id < numberConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", serverAddress+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			logger.Debugf("telling relay: %s", connectionType+"."+codePhrase)

			sendMessage(connectionType+"."+Hash(codePhrase), connection)
			if connectionType == "s" { // this is a sender
				if id == 0 {
					fmt.Println("waiting for other to connect")
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				sendFile(id, connection, codePhrase)
			} else { // this is a receiver
				// receive file
				logger.Debug("receive file")
				fileName, fileIV, fileSalt, fileHash = receiveFile(id, connection, codePhrase)
			}

		}(id)
	}
	wg.Wait()

	if connectionType == "r" {
		catFile(fileName)
		encrypted, err := ioutil.ReadFile(fileName + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		log.Debugf("codePhrase: [%s]", codePhrase)
		log.Debugf("fileSalt: [%s]", fileSalt)
		log.Debugf("fileIV: [%s]", fileIV)
		decrypted, err := Decrypt(encrypted, codePhrase, fileSalt, fileIV)
		if err != nil {
			log.Error(err)
			return
		}
		ioutil.WriteFile(fileName, decrypted, 0644)
		os.Remove(fileName + ".encrypted")
		log.Debugf("\n\n\ndownloaded hash: [%s]", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: [%s]", fileHash)

		if fileHash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", fileName)
		} else {
			fmt.Printf("\nDownloaded %s!", fileName)
		}
	}
}

func catFile(fileNameToReceive string) {
	// cat the file
	os.Remove(fileNameToReceive)
	finished, err := os.Create(fileNameToReceive + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < numberConnections; id++ {
		fh, err := os.Open(fileNameToReceive + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	}

}

func receiveFile(id int, connection net.Conn, codePhrase string) (fileNameToReceive string, iv string, salt string, hashOfFile string) {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for file size")

	bufferFileSize := make([]byte, 10)
	connection.Read(bufferFileSize)
	fileSize, _ := strconv.ParseInt(strings.Trim(string(bufferFileSize), ":"), 10, 64)
	logger.Debugf("filesize: %d", fileSize)

	bufferFileName := make([]byte, 64)
	connection.Read(bufferFileName)
	fileNameToReceive = strings.Trim(string(bufferFileName), ":")
	logger.Debugf("fileName: [%s]", fileNameToReceive)

	ivHex := make([]byte, BUFFERSIZE)
	connection.Read(ivHex)
	iv = strings.Trim(strings.TrimSpace(string(ivHex)), ":")
	iv = strings.Replace(iv, ":", "", -1)
	logger.Debugf("iv: [%s]", iv)

	saltHex := make([]byte, BUFFERSIZE)
	connection.Read(saltHex)
	salt = strings.Trim(strings.TrimSpace(string(saltHex)), ":")
	logger.Debugf("salt: [%s]", salt)

	hashOfFileBytes := make([]byte, BUFFERSIZE)
	connection.Read(hashOfFileBytes)
	hashOfFile = strings.Trim(strings.TrimSpace(string(hashOfFileBytes)), ":")
	logger.Debugf("hashOfFile: [%s]", hashOfFile)

	os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	newFile, err := os.Create(fileNameToReceive + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	if !debugFlag {
		bars[id] = uiprogress.AddBar(int(fileSize)/1024 + 1).AppendCompleted().PrependElapsed()
	}

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		if !debugFlag {
			bars[id].Incr()
		}
		if (fileSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (fileSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-fileSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-fileSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		//Increment the counter
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return
}

func sendFile(id int, connection net.Conn, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	// // Open the file that needs to be send to the client
	// file, err := os.Open(fileName + ".encrypted")
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }
	// defer file.Close()
	// // Get the filename and filesize
	// fileInfo, err := file.Stat()
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }

	numChunks := math.Ceil(float64(len(fileBytes)) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(numberConnections)))

	bytesPerConnection := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == numberConnections {
		bytesPerConnection = int64(len(fileBytes)) - (numberConnections-1)*bytesPerConnection
	}

	if id == 0 || id == numberConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerConnection: %v", bytesPerConnection)
		logger.Debugf("fileNameToSend: %v", path.Base(fileName))
	}

	// send file size
	logger.Debugf("sending fileSize: %d", bytesPerConnection)
	connection.Write([]byte(fillString(strconv.FormatInt(int64(bytesPerConnection), 10), 10)))

	// send fileName
	logger.Debugf("sending fileName: %s", path.Base(fileName))
	connection.Write([]byte(fillString(path.Base(fileName), 64)))

	// send iv
	logger.Debugf("sending iv: %s", fileIV)
	connection.Write([]byte(fillString(fileIV, BUFFERSIZE)))

	// send salt
	logger.Debugf("sending salt: %s", fileSalt)
	connection.Write([]byte(fillString(fileSalt, BUFFERSIZE)))

	// send sha256sum of file
	logger.Debugf("sending sha256sum: %s", fileHash)
	connection.Write([]byte(fillString(fileHash, BUFFERSIZE)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(fileBytes)
	chunkI := 0
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == numberConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<commit_msg>Turn off decryption<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

var bars []*uiprogress.Bar

// runClient spawns threads for parallel uplink/downlink via TCP
func runClient(connectionType string, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"codePhrase": codePhrase,
		"connection": connectionType,
	})
	var wg sync.WaitGroup
	wg.Add(numberConnections)

	uiprogress.Start()
	if !debugFlag {
		bars = make([]*uiprogress.Bar, numberConnections)
	}
	for id := 0; id < numberConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", serverAddress+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			logger.Debugf("telling relay: %s", connectionType+"."+codePhrase)

			sendMessage(connectionType+"."+Hash(codePhrase), connection)
			if connectionType == "s" { // this is a sender
				if id == 0 {
					fmt.Println("waiting for other to connect")
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				sendFile(id, connection, codePhrase)
			} else { // this is a receiver
				// receive file
				logger.Debug("receive file")
				fileName, fileIV, fileSalt, fileHash = receiveFile(id, connection, codePhrase)
			}

		}(id)
	}
	wg.Wait()

	if connectionType == "r" {
		catFile(fileName)
		encrypted, err := ioutil.ReadFile(fileName + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		log.Debugf("codePhrase: [%s]", codePhrase)
		log.Debugf("fileSalt: [%s]", fileSalt)
		log.Debugf("fileIV: [%s]", fileIV)
		decrypted, err := Decrypt(encrypted, codePhrase, fileSalt, fileIV)
		if err != nil {
			log.Error(err)
			return
		}
		ioutil.WriteFile(fileName, decrypted, 0644)
		if !debugFlag {
			os.Remove(fileName + ".encrypted")
		}
		log.Debugf("\n\n\ndownloaded hash: [%s]", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: [%s]", fileHash)

		if fileHash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", fileName)
		} else {
			fmt.Printf("\nDownloaded %s!", fileName)
		}
	}
}

func catFile(fileNameToReceive string) {
	// cat the file
	os.Remove(fileNameToReceive)
	finished, err := os.Create(fileNameToReceive + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < numberConnections; id++ {
		fh, err := os.Open(fileNameToReceive + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	}

}

func receiveFile(id int, connection net.Conn, codePhrase string) (fileNameToReceive string, iv string, salt string, hashOfFile string) {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for file size")

	bufferFileSize := make([]byte, 10)
	connection.Read(bufferFileSize)
	fileSize, _ := strconv.ParseInt(strings.Trim(string(bufferFileSize), ":"), 10, 64)
	logger.Debugf("filesize: %d", fileSize)

	bufferFileName := make([]byte, 64)
	connection.Read(bufferFileName)
	fileNameToReceive = strings.Trim(string(bufferFileName), ":")
	logger.Debugf("fileName: [%s]", fileNameToReceive)

	ivHex := make([]byte, BUFFERSIZE)
	connection.Read(ivHex)
	iv = strings.Trim(strings.TrimSpace(string(ivHex)), ":")
	iv = strings.Replace(iv, ":", "", -1)
	logger.Debugf("iv: [%s]", iv)

	saltHex := make([]byte, BUFFERSIZE)
	connection.Read(saltHex)
	salt = strings.Trim(strings.TrimSpace(string(saltHex)), ":")
	logger.Debugf("salt: [%s]", salt)

	hashOfFileBytes := make([]byte, BUFFERSIZE)
	connection.Read(hashOfFileBytes)
	hashOfFile = strings.Trim(strings.TrimSpace(string(hashOfFileBytes)), ":")
	logger.Debugf("hashOfFile: [%s]", hashOfFile)

	os.Remove(fileNameToReceive + "." + strconv.Itoa(id))
	newFile, err := os.Create(fileNameToReceive + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	if !debugFlag {
		bars[id] = uiprogress.AddBar(int(fileSize)/1024 + 1).AppendCompleted().PrependElapsed()
	}

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		if !debugFlag {
			bars[id].Incr()
		}
		if (fileSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (fileSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-fileSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-fileSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		//Increment the counter
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return
}

func sendFile(id int, connection net.Conn, codePhrase string) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	// // Open the file that needs to be send to the client
	// file, err := os.Open(fileName + ".encrypted")
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }
	// defer file.Close()
	// // Get the filename and filesize
	// fileInfo, err := file.Stat()
	// if err != nil {
	// 	logger.Error(err)
	// 	return
	// }

	numChunks := math.Ceil(float64(len(fileBytes)) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(numberConnections)))

	bytesPerConnection := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == numberConnections {
		bytesPerConnection = int64(len(fileBytes)) - (numberConnections-1)*bytesPerConnection
	}

	if id == 0 || id == numberConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerConnection: %v", bytesPerConnection)
		logger.Debugf("fileNameToSend: %v", path.Base(fileName))
	}

	// send file size
	logger.Debugf("sending fileSize: %d", bytesPerConnection)
	connection.Write([]byte(fillString(strconv.FormatInt(int64(bytesPerConnection), 10), 10)))

	// send fileName
	logger.Debugf("sending fileName: %s", path.Base(fileName))
	connection.Write([]byte(fillString(path.Base(fileName), 64)))

	// send iv
	logger.Debugf("sending iv: %s", fileIV)
	connection.Write([]byte(fillString(fileIV, BUFFERSIZE)))

	// send salt
	logger.Debugf("sending salt: %s", fileSalt)
	connection.Write([]byte(fillString(fileSalt, BUFFERSIZE)))

	// send sha256sum of file
	logger.Debugf("sending sha256sum: %s", fileHash)
	connection.Write([]byte(fillString(fileHash, BUFFERSIZE)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(fileBytes)
	chunkI := 0
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == numberConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/joshsoftware/curem/config"
	"labix.org/v2/mgo/bson"
)

// contact type holds the fields related to a particular contact.
// omitempty tag will make sure the database doesn't contain content like:
//
//   {
//    _id: someId
//    company: ABC
//    Person: Xyz
//    Phone:
//    Skype:
//    Country:
//   }
// Instead, it will store the above data as:
//
//   {
//    _id: someId
//    company: ABC
//    Person: Xyz
//   }
type contact struct {
	Id      bson.ObjectId `bson:"_id"               json:"id"`
	Company string        `bson:"company,omitempty" json:"company,omitempty"`
	Person  string        `bson:"person,omitempty"  json:"person,omitempty"`
	Email   string        `bson:"email,omitempty"   json:"email,omitempty"`
	Phone   string        `bson:"phone,omitempty"   json:"phone,omitempty"`
	SkypeId string        `bson:"skypeid,omitempty" json:"skypeid,omitempty"`
	Country string        `bson:"country,omitempty" json:"country,omitempty"`
}

// NewContact takes the fields of a contact, initializes a struct of contact type and returns
// the pointer to that struct.
// Also, It inserts the contact data into a mongoDB collection, which is passed as the first parameter.
func NewContact(company, person, email, phone, skypeid, country string) (*contact, error) {
	collection := config.Db.C("newcontact")
	doc := contact{
		Id:      bson.NewObjectId(),
		Company: company,
		Person:  person,
		Email:   email,
		Phone:   phone,
		SkypeId: skypeid,
		Country: country,
	}
	err := collection.Insert(doc)
	if err != nil {
		return &contact{}, err
	}
	return &doc, nil
}

func GetContact(i bson.ObjectId) (*contact, error) {
	collection := config.Db.C("newcontact")
	var c contact
	err := collection.FindId(i).One(&c)
	if err != nil {
		return &contact{}, err
	}
	return &c, nil
}

func (c *contact) Delete() error {
	i := c.Id
	collection := config.Db.C("newcontact")
	err := collection.RemoveId(i)
	return err
}
<commit_msg>Add update method and use config variables wherever necessary<commit_after>package main

import (
	"github.com/joshsoftware/curem/config"
	"labix.org/v2/mgo/bson"
)

// contact type holds the fields related to a particular contact.
// omitempty tag will make sure the database doesn't contain content like:
//
//   {
//    _id: someId
//    company: ABC
//    Person: Xyz
//    Phone:
//    Skype:
//    Country:
//   }
// Instead, it will store the above data as:
//
//   {
//    _id: someId
//    company: ABC
//    Person: Xyz
//   }
type contact struct {
	Id      bson.ObjectId `bson:"_id"               json:"id"`
	Company string        `bson:"company,omitempty" json:"company,omitempty"`
	Person  string        `bson:"person,omitempty"  json:"person,omitempty"`
	Email   string        `bson:"email,omitempty"   json:"email,omitempty"`
	Phone   string        `bson:"phone,omitempty"   json:"phone,omitempty"`
	SkypeId string        `bson:"skypeid,omitempty" json:"skypeid,omitempty"`
	Country string        `bson:"country,omitempty" json:"country,omitempty"`
}

// NewContact takes the fields of a contact, initializes a struct of contact type and returns
// the pointer to that struct.
func NewContact(company, person, email, phone, skypeid, country string) (*contact, error) {
	doc := contact{
		Id:      bson.NewObjectId(),
		Company: company,
		Person:  person,
		Email:   email,
		Phone:   phone,
		SkypeId: skypeid,
		Country: country,
	}
	err := config.ContactsCollection.Insert(doc)
	if err != nil {
		return &contact{}, err
	}
	return &doc, nil
}

func GetContact(i bson.ObjectId) (*contact, error) {
	var c contact
	err := config.ContactsCollection.FindId(i).One(&c)
	if err != nil {
		return &contact{}, err
	}
	return &c, nil
}

// Update updates the contact in the database.
// First, fetch a contact from the database and change the necessary fields.
// Then call the Update method on the contact object.
func (c *contact) Update() error {
	_, err := config.ContactsCollection.UpsertId(c.Id, c)
	return err
}

// Delete deletes the contact from the database.
func (c *contact) Delete() error {
	return config.ContactsCollection.RemoveId(c.Id)
}
<|endoftext|>


<commit_before>package request

import (
	"encoding/json"
	"github.com/MatthewHartstonge/storage/client"
	"github.com/ory/fosite"
	"github.com/pkg/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

const (
	mongoCollectionOpenIDSessions    = "OpenIDConnectSessions"
	mongoCollectionAccessTokens      = "AccessTokens"
	mongoCollectionRefreshTokens     = "RefreshTokens"
	mongoCollectionAuthorizationCode = "AuthorizationCode"
)

// MongoManager manages the main Mongo Session for a Request.
type MongoManager struct {
	client.MongoManager

	// DB is the Mongo connection that holds the base session that can be copied and closed.
	DB *mgo.Database

	// TODO: Add AES cipher for Token Encryption?
}

// Given a request from fosite, marshals to a form that enables storing the request in mongo
func mongoCollectionFromRequest(signature string, r fosite.Requester) (*mongoRequestData, error) {
	session, err := json.Marshal(r.GetSession())
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &mongoRequestData{
		ID:            r.GetID(),
		RequestedAt:   r.GetRequestedAt(),
		Signature:     signature,
		ClientID:      r.GetClient().GetID(),
		Scopes:        r.GetRequestedScopes(),
		GrantedScopes: r.GetGrantedScopes(),
		Form:          r.GetRequestForm().Encode(),
		Session:       session,
	}, nil

}

// createSession stores a session to a specific mongo collection
func (m *MongoManager) createSession(signature string, requester fosite.Requester, collectionName string) error {
	data, err := mongoCollectionFromRequest(signature, requester)
	if err != nil {
		return err
	}

	c := m.DB.C(collectionName).With(m.DB.Session.Copy())
	defer c.Database.Session.Close()
	if err := m.DB.C(collectionName).Insert(data); err != nil {
		return errors.WithStack(err)
	}
	return nil
}

// findSessionBySignature finds a session within a specific mongo collection
func (m *MongoManager) findSessionBySignature(signature string, session fosite.Session, collectionName string) (fosite.Requester, error) {
	var d *mongoRequestData
	c := m.DB.C(collectionName).With(m.DB.Session.Copy())
	defer c.Database.Session.Close()
	if err := m.DB.C(collectionName).Find(bson.M{"signature": signature}).One(d); err == mgo.ErrNotFound {
		return nil, fosite.ErrNotFound
	} else if err != nil {
		return nil, errors.WithStack(err)
	}

	return d.toRequest(session, m.MongoManager)
}
<commit_msg>:zap: request: Use the returned collection to access the copied session<commit_after>package request

import (
	"encoding/json"
	"github.com/MatthewHartstonge/storage/client"
	"github.com/ory/fosite"
	"github.com/pkg/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

const (
	mongoCollectionOpenIDSessions    = "OpenIDConnectSessions"
	mongoCollectionAccessTokens      = "AccessTokens"
	mongoCollectionRefreshTokens     = "RefreshTokens"
	mongoCollectionAuthorizationCode = "AuthorizationCode"
)

// MongoManager manages the main Mongo Session for a Request.
type MongoManager struct {
	client.MongoManager

	// DB is the Mongo connection that holds the base session that can be copied and closed.
	DB *mgo.Database

	// TODO: Add AES cipher for Token Encryption?
}

// Given a request from fosite, marshals to a form that enables storing the request in mongo
func mongoCollectionFromRequest(signature string, r fosite.Requester) (*mongoRequestData, error) {
	session, err := json.Marshal(r.GetSession())
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &mongoRequestData{
		ID:            r.GetID(),
		RequestedAt:   r.GetRequestedAt(),
		Signature:     signature,
		ClientID:      r.GetClient().GetID(),
		Scopes:        r.GetRequestedScopes(),
		GrantedScopes: r.GetGrantedScopes(),
		Form:          r.GetRequestForm().Encode(),
		Session:       session,
	}, nil

}

// createSession stores a session to a specific mongo collection
func (m *MongoManager) createSession(signature string, requester fosite.Requester, collectionName string) error {
	data, err := mongoCollectionFromRequest(signature, requester)
	if err != nil {
		return err
	}

	c := m.DB.C(collectionName).With(m.DB.Session.Copy())
	defer c.Database.Session.Close()
	if err := c.Insert(data); err != nil {
		return errors.WithStack(err)
	}
	return nil
}

// findSessionBySignature finds a session within a specific mongo collection
func (m *MongoManager) findSessionBySignature(signature string, session fosite.Session, collectionName string) (fosite.Requester, error) {
	var d *mongoRequestData
	c := m.DB.C(collectionName).With(m.DB.Session.Copy())
	defer c.Database.Session.Close()
	if err := c.Find(bson.M{"signature": signature}).One(d); err == mgo.ErrNotFound {
		return nil, fosite.ErrNotFound
	} else if err != nil {
		return nil, errors.WithStack(err)
	}

	return d.toRequest(session, m.MongoManager)
}
<|endoftext|>


<commit_before>package upkeep

import (
	pxy "github.com/jeffjen/ambd/ambctl/arg"

	log "github.com/Sirupsen/logrus"
	docker "github.com/fsouza/go-dockerclient"
	dri "github.com/jeffjen/docker-monitor/driver"
	disc "github.com/jeffjen/go-discovery"
	node "github.com/jeffjen/go-discovery/info"

	etcd "github.com/coreos/etcd/client"
	ctx "golang.org/x/net/context"

	"fmt"
	"os"
	"path"
	"time"
)

func init() {
	var level = os.Getenv("LOG_LEVEL")
	switch level {
	case "DEBUG":
		log.SetLevel(log.DebugLevel)
		break
	case "INFO":
		log.SetLevel(log.InfoLevel)
		break
	case "WARNING":
		log.SetLevel(log.WarnLevel)
		break
	case "ERROR":
		log.SetLevel(log.ErrorLevel)
		break
	case "FATAL":
		log.SetLevel(log.FatalLevel)
		break
	case "PANIC":
		log.SetLevel(log.PanicLevel)
		break
	default:
		log.SetLevel(log.InfoLevel)
		break
	}
}

const (
	UpkeepTimeout = 3 * time.Second

	ProbeTimeout = 100 * time.Millisecond

	MaxFailAttemps = 3
)

const (
	ServiceUp          = "up"
	ServiceUnavailable = "unavailable"
	ServiceDown        = "down"
	ServiceRemoved     = "die"
)

type Service struct {
	State string `json:"State"`

	Hb  time.Duration `json: "Heartbeat"`
	TTL time.Duration `json: "TTL"`

	PHb           time.Duration `json: "ProbeHeartbeat"`
	ProbeType     string        `json: "ProbeType"`
	ProbeEndpoint string        `json: "ProbeEndpoint"`

	Id       string           `json: "ContainerID"`
	Srv      string           `json: "Service"`
	Port     string           `json: "Port"`
	Net      []docker.APIPort `json: "Net"`
	Proxy    []pxy.Info       `json: "Proxy"`
	ProxyCfg string           `json: "ProxyCfg"`

	Key []string `json: "Key"`
}

func pushReport(serv *srv) {
	pushlog := log.WithFields(log.Fields{"srv": serv.Srv, "state": serv.State})
	go func() {
		wk, cancel := ctx.WithTimeout(ctx.Background(), 5*time.Second)
		defer cancel()
		if err := report.Push(wk, serv); err != nil {
			pushlog.WithFields(log.Fields{"err": err}).Warning("push")
		}
	}()
}

type srv struct {
	*Service // embedded type

	kAPI etcd.KeysAPI
	opts *etcd.SetOptions

	driver dri.Driver
}

func (serv *srv) keep(c ctx.Context) (err error) {
	for _, k := range serv.Key {
		_, err = serv.kAPI.Set(c, k, node.MetaData, serv.opts)
		if err != nil {
			break
		}
	}
	preState := serv.State
	if err != nil {
		serv.State = ServiceUnavailable
		serv.opts.PrevExist = etcd.PrevIgnore
		pushReport(serv)
	} else {
		serv.State = ServiceUp
		serv.opts.PrevExist = etcd.PrevExist
		if preState != ServiceUp {
			pushReport(serv)
		}
	}
	return
}

func (serv *srv) probe(c ctx.Context) (err error) {
	err = serv.driver.Probe(c)
	log.WithFields(log.Fields{"err": err, "srv": serv.Srv}).Debug("probe")
	return
}

func (serv *srv) urk() {
	serv.State = ServiceUnavailable
	serv.opts.PrevExist = etcd.PrevIgnore
	pushReport(serv)
}

func (serv *srv) down() {
	serv.State = ServiceDown
	serv.opts.PrevExist = etcd.PrevIgnore
	pushReport(serv)
}

func (serv *srv) die() {
	serv.State = ServiceRemoved
	pushReport(serv)
}

func Get(iden string) (s *Service) {
	s, _ = rec.Get(iden).(*Service)
	return
}

func Place(service *Service) {
	if r, ok := Record[service.Id]; ok {
		r.Abort()
	}
	rec.Set(service.Id, service)
	Register(service)
}

func Register(service *Service) {
	alloc := AllocHelper(service.ProbeType)
	service.Key = make([]string, 0)

	serv := &srv{
		service,
		etcd.NewKeysAPI(disc.NewDiscovery()),
		&etcd.SetOptions{TTL: service.TTL, PrevExist: etcd.PrevIgnore},
		nil,
	}

	logger := log.WithFields(log.Fields{
		"ID": serv.Id[:12], "srv": serv.Srv, "heartbeat": serv.Hb, "ttl": serv.TTL,
	})

	// Advertise Key on the discovery service
	if serv.Port != "" {
		serv.Key = append(serv.Key, fmt.Sprintf("%s/%s:%s", serv.Srv, Advertise, serv.Port))
	} else if len(serv.Net) > 0 {
		serv.Key = make([]string, 0)
		for _, p := range serv.Net {
			if p.PublicPort != 0 && p.IP == "0.0.0.0" {
				serv.Key = append(serv.Key, fmt.Sprintf("%s/%s:%d", serv.Srv, Advertise, p.PublicPort))
			}
		}
	}

	var endpoint string
	if serv.ProbeEndpoint == "" {
		endpoint = path.Base(serv.Key[0])
	} else {
		endpoint = fmt.Sprintf("%s/%s", path.Base(serv.Key[0]), serv.ProbeEndpoint)
	}

	// TODO:  setup driver for probing
	driver, drr := alloc(endpoint)
	if drr != nil {
		logger.WithFields(log.Fields{"err": drr}).Error("-register")
		return
	}
	serv.driver = driver
	logger.Debug("+register")

	wk, abort := ctx.WithCancel(RootContext)
	go func() {
		defer serv.driver.Close()

		// Request to establish proxy port to ambassador
		openProxyConfig(serv.ProxyCfg, serv.Proxy)

		// setup work cycle
		heartbeat, probe := time.NewTicker(serv.Hb), time.NewTicker(serv.PHb)
		defer func() {
			heartbeat.Stop()
			probe.Stop()
		}()

		logger.Info("start")
		func() {
			d, abort := ctx.WithTimeout(wk, UpkeepTimeout)
			if err := serv.keep(d); err != nil {
				logger.WithFields(log.Fields{"err": err, "state": serv.opts.PrevExist}).Error("-up")
			} else {
				logger.Info("+up")
			}
			abort() // release
		}()

		var chk = NewFail(MaxFailAttemps)
		for yay := true; yay; {
			select {
			case <-heartbeat.C:
				if !chk.Pass() {
					serv.urk()
					logger.Error("!up")
				} else {
					d, abort := ctx.WithTimeout(wk, UpkeepTimeout)
					if err := serv.keep(d); err != nil {
						logger.WithFields(log.Fields{"err": err, "state": serv.opts.PrevExist}).Error("-up")
					} else {
						logger.Info("+up")
					}
					abort() // release
				}

			case <-probe.C:
				d, abort := ctx.WithTimeout(wk, ProbeTimeout)
				if err := serv.probe(d); err != nil {
					count := chk.Bad()
					logger.WithFields(log.Fields{"err": err, "fail": count}).Warning("-probe")
				} else {
					chk.Good()
					logger.Debug("+probe")
				}
				abort() // release

			case <-wk.Done():
				serv.down()
				logger.Warning("down")
				yay = false
			}
		}
	}()

	Record[serv.Id] = &RunningRecord{serv.Srv, abort} // register abort function for this service
}

func Suspend(iden string) {
	if r, ok := Record[iden]; ok {
		r.Abort()
	}
}

func Unregister(iden string) {
	if r, ok := Record[iden]; ok {
		delete(Record, iden)
		r.Abort()
		go func() {
			serv := &srv{rec.Get(iden).(*Service), nil, nil, nil}
			serv.die()
			rec.Del(iden)
		}()
		log.WithFields(log.Fields{"ID": iden[:12], "srv": r.Srv}).Warning("die")
	}
}
<commit_msg>TMP: FIX: extend ProbeTimeout<commit_after>package upkeep

import (
	pxy "github.com/jeffjen/ambd/ambctl/arg"

	log "github.com/Sirupsen/logrus"
	docker "github.com/fsouza/go-dockerclient"
	dri "github.com/jeffjen/docker-monitor/driver"
	disc "github.com/jeffjen/go-discovery"
	node "github.com/jeffjen/go-discovery/info"

	etcd "github.com/coreos/etcd/client"
	ctx "golang.org/x/net/context"

	"fmt"
	"os"
	"path"
	"time"
)

func init() {
	var level = os.Getenv("LOG_LEVEL")
	switch level {
	case "DEBUG":
		log.SetLevel(log.DebugLevel)
		break
	case "INFO":
		log.SetLevel(log.InfoLevel)
		break
	case "WARNING":
		log.SetLevel(log.WarnLevel)
		break
	case "ERROR":
		log.SetLevel(log.ErrorLevel)
		break
	case "FATAL":
		log.SetLevel(log.FatalLevel)
		break
	case "PANIC":
		log.SetLevel(log.PanicLevel)
		break
	default:
		log.SetLevel(log.InfoLevel)
		break
	}
}

const (
	UpkeepTimeout = 3 * time.Second

	ProbeTimeout = 1 * time.Second

	MaxFailAttemps = 3
)

const (
	ServiceUp          = "up"
	ServiceUnavailable = "unavailable"
	ServiceDown        = "down"
	ServiceRemoved     = "die"
)

type Service struct {
	State string `json:"State"`

	Hb  time.Duration `json: "Heartbeat"`
	TTL time.Duration `json: "TTL"`

	PHb           time.Duration `json: "ProbeHeartbeat"`
	ProbeType     string        `json: "ProbeType"`
	ProbeEndpoint string        `json: "ProbeEndpoint"`

	Id       string           `json: "ContainerID"`
	Srv      string           `json: "Service"`
	Port     string           `json: "Port"`
	Net      []docker.APIPort `json: "Net"`
	Proxy    []pxy.Info       `json: "Proxy"`
	ProxyCfg string           `json: "ProxyCfg"`

	Key []string `json: "Key"`
}

func pushReport(serv *srv) {
	pushlog := log.WithFields(log.Fields{"srv": serv.Srv, "state": serv.State})
	go func() {
		wk, cancel := ctx.WithTimeout(ctx.Background(), 5*time.Second)
		defer cancel()
		if err := report.Push(wk, serv); err != nil {
			pushlog.WithFields(log.Fields{"err": err}).Warning("push")
		}
	}()
}

type srv struct {
	*Service // embedded type

	kAPI etcd.KeysAPI
	opts *etcd.SetOptions

	driver dri.Driver
}

func (serv *srv) keep(c ctx.Context) (err error) {
	for _, k := range serv.Key {
		_, err = serv.kAPI.Set(c, k, node.MetaData, serv.opts)
		if err != nil {
			break
		}
	}
	preState := serv.State
	if err != nil {
		serv.State = ServiceUnavailable
		serv.opts.PrevExist = etcd.PrevIgnore
		pushReport(serv)
	} else {
		serv.State = ServiceUp
		serv.opts.PrevExist = etcd.PrevExist
		if preState != ServiceUp {
			pushReport(serv)
		}
	}
	return
}

func (serv *srv) probe(c ctx.Context) (err error) {
	err = serv.driver.Probe(c)
	log.WithFields(log.Fields{"err": err, "srv": serv.Srv}).Debug("probe")
	return
}

func (serv *srv) urk() {
	serv.State = ServiceUnavailable
	serv.opts.PrevExist = etcd.PrevIgnore
	pushReport(serv)
}

func (serv *srv) down() {
	serv.State = ServiceDown
	serv.opts.PrevExist = etcd.PrevIgnore
	pushReport(serv)
}

func (serv *srv) die() {
	serv.State = ServiceRemoved
	pushReport(serv)
}

func Get(iden string) (s *Service) {
	s, _ = rec.Get(iden).(*Service)
	return
}

func Place(service *Service) {
	if r, ok := Record[service.Id]; ok {
		r.Abort()
	}
	rec.Set(service.Id, service)
	Register(service)
}

func Register(service *Service) {
	alloc := AllocHelper(service.ProbeType)
	service.Key = make([]string, 0)

	serv := &srv{
		service,
		etcd.NewKeysAPI(disc.NewDiscovery()),
		&etcd.SetOptions{TTL: service.TTL, PrevExist: etcd.PrevIgnore},
		nil,
	}

	logger := log.WithFields(log.Fields{
		"ID": serv.Id[:12], "srv": serv.Srv, "heartbeat": serv.Hb, "ttl": serv.TTL,
	})

	// Advertise Key on the discovery service
	if serv.Port != "" {
		serv.Key = append(serv.Key, fmt.Sprintf("%s/%s:%s", serv.Srv, Advertise, serv.Port))
	} else if len(serv.Net) > 0 {
		serv.Key = make([]string, 0)
		for _, p := range serv.Net {
			if p.PublicPort != 0 && p.IP == "0.0.0.0" {
				serv.Key = append(serv.Key, fmt.Sprintf("%s/%s:%d", serv.Srv, Advertise, p.PublicPort))
			}
		}
	}

	var endpoint string
	if serv.ProbeEndpoint == "" {
		endpoint = path.Base(serv.Key[0])
	} else {
		endpoint = fmt.Sprintf("%s/%s", path.Base(serv.Key[0]), serv.ProbeEndpoint)
	}

	// TODO:  setup driver for probing
	driver, drr := alloc(endpoint)
	if drr != nil {
		logger.WithFields(log.Fields{"err": drr}).Error("-register")
		return
	}
	serv.driver = driver
	logger.Debug("+register")

	wk, abort := ctx.WithCancel(RootContext)
	go func() {
		defer serv.driver.Close()

		// Request to establish proxy port to ambassador
		openProxyConfig(serv.ProxyCfg, serv.Proxy)

		// setup work cycle
		heartbeat, probe := time.NewTicker(serv.Hb), time.NewTicker(serv.PHb)
		defer func() {
			heartbeat.Stop()
			probe.Stop()
		}()

		logger.Info("start")
		func() {
			d, abort := ctx.WithTimeout(wk, UpkeepTimeout)
			if err := serv.keep(d); err != nil {
				logger.WithFields(log.Fields{"err": err, "state": serv.opts.PrevExist}).Error("-up")
			} else {
				logger.Info("+up")
			}
			abort() // release
		}()

		var chk = NewFail(MaxFailAttemps)
		for yay := true; yay; {
			select {
			case <-heartbeat.C:
				if !chk.Pass() {
					serv.urk()
					logger.Error("!up")
				} else {
					d, abort := ctx.WithTimeout(wk, UpkeepTimeout)
					if err := serv.keep(d); err != nil {
						logger.WithFields(log.Fields{"err": err, "state": serv.opts.PrevExist}).Error("-up")
					} else {
						logger.Info("+up")
					}
					abort() // release
				}

			case <-probe.C:
				d, abort := ctx.WithTimeout(wk, ProbeTimeout)
				if err := serv.probe(d); err != nil {
					count := chk.Bad()
					logger.WithFields(log.Fields{"err": err, "fail": count}).Warning("-probe")
				} else {
					chk.Good()
					logger.Debug("+probe")
				}
				abort() // release

			case <-wk.Done():
				serv.down()
				logger.Warning("down")
				yay = false
			}
		}
	}()

	Record[serv.Id] = &RunningRecord{serv.Srv, abort} // register abort function for this service
}

func Suspend(iden string) {
	if r, ok := Record[iden]; ok {
		r.Abort()
	}
}

func Unregister(iden string) {
	if r, ok := Record[iden]; ok {
		delete(Record, iden)
		r.Abort()
		go func() {
			serv := &srv{rec.Get(iden).(*Service), nil, nil, nil}
			serv.die()
			rec.Del(iden)
		}()
		log.WithFields(log.Fields{"ID": iden[:12], "srv": r.Srv}).Warning("die")
	}
}
<|endoftext|>


<commit_before>package upstream

import (
	"fmt"
	"github.com/ngmoco/falcore"
	"net"
	"net/http"
	"time"
	"io"
	"bytes"
)

type passThruReadCloser struct {
	io.Reader
	io.Closer
}

type Upstream struct {
	// The upstream host to connect to
	Host string
	// The port on the upstream host
	Port int
	// Default 60 seconds
	Timeout time.Duration
	// Will ignore https on the incoming request and always upstream http
	ForceHttp bool
	// Ping URL Path-only for checking upness
	PingPath string

	transport *http.Transport
	host      string
	tcpaddr   *net.TCPAddr
	tcpconn   *net.TCPConn
}

func NewUpstream(host string, port int, forceHttp bool) *Upstream {
	u := new(Upstream)
	u.Host = host
	u.Port = port
	u.ForceHttp = forceHttp
	ips, err := net.LookupIP(host)
	var ip net.IP = nil
	for i := range ips {
		ip = ips[i].To4()
		if ip != nil {
			break
		}
	}
	if err == nil && ip != nil {
		u.tcpaddr = new(net.TCPAddr)
		u.tcpaddr.Port = port
		u.tcpaddr.IP = ip
	} else {
		falcore.Warn("Can't get IP addr for %v: %v", host, err)
	}
	u.Timeout = 60e9
	u.host = fmt.Sprintf("%v:%v", u.Host, u.Port)

	u.transport = new(http.Transport)

	u.transport.Dial = func(n, addr string) (c net.Conn, err error) {
		falcore.Fine("Dialing connection to %v", u.tcpaddr)
		var ctcp *net.TCPConn
		ctcp, err = net.DialTCP("tcp4", nil, u.tcpaddr)
		if ctcp != nil {
			u.tcpconn = ctcp
			u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
		}
		if err != nil {
			falcore.Error("Dial Failed: %v", err)
		}
		return ctcp, err
	}
	u.transport.MaxIdleConnsPerHost = 15
	return u
}

// Alter the number of connections to multiplex with
func (u *Upstream) SetPoolSize(size int) {
	u.transport.MaxIdleConnsPerHost = size
}

func (u *Upstream) FilterRequest(request *falcore.Request) (res *http.Response) {
	var err error
	req := request.HttpRequest

	// Force the upstream to use http
	if u.ForceHttp || req.URL.Scheme == "" {
		req.URL.Scheme = "http"
		req.URL.Host = req.Host
	}
	before := time.Now()
	req.Header.Set("Connection", "Keep-Alive")
	if u.tcpconn != nil {
		u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
	}
	var upstrRes *http.Response
	upstrRes, err = u.transport.RoundTrip(req)
	diff := falcore.TimeDiff(before, time.Now())
	if err == nil {
		// Copy response over to new record.  Remove connection noise.  Add some sanity.
		res = falcore.SimpleResponse(req, upstrRes.StatusCode, nil, "")
		if upstrRes.ContentLength > 0 && upstrRes.Body != nil {
			res.Body = upstrRes.Body
		} else if upstrRes.ContentLength == 0 && upstrRes.Body != nil {
			// Any bytes?
			var testBuf [1]byte
			n, _ := io.ReadFull(upstrRes.Body, testBuf[:])
			if n == 1 {
				// Yes there are.  Chunked it is.
				res.TransferEncoding = []string{"chunked"}
				res.ContentLength = -1
				rc := &passThruReadCloser{
					io.MultiReader(bytes.NewBuffer(testBuf[:]), upstrRes.Body),
					upstrRes.Body,
				}

				res.Body = rc
			}
		} else if upstrRes.Body != nil {
			res.Body = upstrRes.Body
			res.ContentLength = -1
			res.TransferEncoding = []string{"chunked"}
		}
		// Copy over headers with a few exceptions
		res.Header = make(http.Header)
		for hn, hv := range upstrRes.Header {
			switch hn {
			case "Content-Length":
			case "Connection":
			case "Transfer-Encoding":
			default:
				res.Header[hn] = hv
			}
		}
	} else {
		if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
			falcore.Error("%s Upstream Timeout error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 504, nil, "Gateway Timeout\n")
			request.CurrentStage.Status = 2 // Fail
		} else {
			falcore.Error("%s Upstream error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 502, nil, "Bad Gateway\n")
			request.CurrentStage.Status = 2 // Fail
		}
	}
	falcore.Debug("%s [%s] [%s] %s s=%d Time=%.4f", request.ID, req.Method, u.host, req.URL, res.StatusCode, diff)
	return
}

func (u *Upstream) ping() (up bool, ok bool) {
	if u.PingPath != "" {
		// the url must be syntactically valid for this to work but the host will be ignored because we
		// are overriding the connection always
		request, err := http.NewRequest("GET", "http://localhost"+u.PingPath, nil)
		request.Header.Set("Connection", "Keep-Alive") // not sure if this should be here for a ping
		if err != nil {
			falcore.Error("Bad Ping request: %v", err)
			return false, true
		}
		if u.tcpconn != nil {
			u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
		}
		res, err := u.transport.RoundTrip(request)

		if err != nil {
			falcore.Error("Failed Ping to %v:%v: %v", u.Host, u.Port, err)
			return false, true
		} else {
			res.Body.Close()
		}
		if res.StatusCode == 200 {
			return true, true
		}
		falcore.Error("Failed Ping to %v:%v: %v", u.Host, u.Port, res.Status)
		// bad status
		return false, true
	}
	return false, false
}
<commit_msg>don't forget the content length<commit_after>package upstream

import (
	"fmt"
	"github.com/ngmoco/falcore"
	"net"
	"net/http"
	"time"
	"io"
	"bytes"
)

type passThruReadCloser struct {
	io.Reader
	io.Closer
}

type Upstream struct {
	// The upstream host to connect to
	Host string
	// The port on the upstream host
	Port int
	// Default 60 seconds
	Timeout time.Duration
	// Will ignore https on the incoming request and always upstream http
	ForceHttp bool
	// Ping URL Path-only for checking upness
	PingPath string

	transport *http.Transport
	host      string
	tcpaddr   *net.TCPAddr
	tcpconn   *net.TCPConn
}

func NewUpstream(host string, port int, forceHttp bool) *Upstream {
	u := new(Upstream)
	u.Host = host
	u.Port = port
	u.ForceHttp = forceHttp
	ips, err := net.LookupIP(host)
	var ip net.IP = nil
	for i := range ips {
		ip = ips[i].To4()
		if ip != nil {
			break
		}
	}
	if err == nil && ip != nil {
		u.tcpaddr = new(net.TCPAddr)
		u.tcpaddr.Port = port
		u.tcpaddr.IP = ip
	} else {
		falcore.Warn("Can't get IP addr for %v: %v", host, err)
	}
	u.Timeout = 60e9
	u.host = fmt.Sprintf("%v:%v", u.Host, u.Port)

	u.transport = new(http.Transport)

	u.transport.Dial = func(n, addr string) (c net.Conn, err error) {
		falcore.Fine("Dialing connection to %v", u.tcpaddr)
		var ctcp *net.TCPConn
		ctcp, err = net.DialTCP("tcp4", nil, u.tcpaddr)
		if ctcp != nil {
			u.tcpconn = ctcp
			u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
		}
		if err != nil {
			falcore.Error("Dial Failed: %v", err)
		}
		return ctcp, err
	}
	u.transport.MaxIdleConnsPerHost = 15
	return u
}

// Alter the number of connections to multiplex with
func (u *Upstream) SetPoolSize(size int) {
	u.transport.MaxIdleConnsPerHost = size
}

func (u *Upstream) FilterRequest(request *falcore.Request) (res *http.Response) {
	var err error
	req := request.HttpRequest

	// Force the upstream to use http
	if u.ForceHttp || req.URL.Scheme == "" {
		req.URL.Scheme = "http"
		req.URL.Host = req.Host
	}
	before := time.Now()
	req.Header.Set("Connection", "Keep-Alive")
	if u.tcpconn != nil {
		u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
	}
	var upstrRes *http.Response
	upstrRes, err = u.transport.RoundTrip(req)
	diff := falcore.TimeDiff(before, time.Now())
	if err == nil {
		// Copy response over to new record.  Remove connection noise.  Add some sanity.
		res = falcore.SimpleResponse(req, upstrRes.StatusCode, nil, "")
		if upstrRes.ContentLength > 0 && upstrRes.Body != nil {
			res.ContentLength = upstrRes.ContentLength
			res.Body = upstrRes.Body
		} else if upstrRes.ContentLength == 0 && upstrRes.Body != nil {
			// Any bytes?
			var testBuf [1]byte
			n, _ := io.ReadFull(upstrRes.Body, testBuf[:])
			if n == 1 {
				// Yes there are.  Chunked it is.
				res.TransferEncoding = []string{"chunked"}
				res.ContentLength = -1
				rc := &passThruReadCloser{
					io.MultiReader(bytes.NewBuffer(testBuf[:]), upstrRes.Body),
					upstrRes.Body,
				}

				res.Body = rc
			}
		} else if upstrRes.Body != nil {
			res.Body = upstrRes.Body
			res.ContentLength = -1
			res.TransferEncoding = []string{"chunked"}
		}
		// Copy over headers with a few exceptions
		res.Header = make(http.Header)
		for hn, hv := range upstrRes.Header {
			switch hn {
			case "Content-Length":
			case "Connection":
			case "Transfer-Encoding":
			default:
				res.Header[hn] = hv
			}
		}
	} else {
		if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
			falcore.Error("%s Upstream Timeout error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 504, nil, "Gateway Timeout\n")
			request.CurrentStage.Status = 2 // Fail
		} else {
			falcore.Error("%s Upstream error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 502, nil, "Bad Gateway\n")
			request.CurrentStage.Status = 2 // Fail
		}
	}
	falcore.Debug("%s [%s] [%s] %s s=%d Time=%.4f", request.ID, req.Method, u.host, req.URL, res.StatusCode, diff)
	return
}

func (u *Upstream) ping() (up bool, ok bool) {
	if u.PingPath != "" {
		// the url must be syntactically valid for this to work but the host will be ignored because we
		// are overriding the connection always
		request, err := http.NewRequest("GET", "http://localhost"+u.PingPath, nil)
		request.Header.Set("Connection", "Keep-Alive") // not sure if this should be here for a ping
		if err != nil {
			falcore.Error("Bad Ping request: %v", err)
			return false, true
		}
		if u.tcpconn != nil {
			u.tcpconn.SetDeadline(time.Now().Add(u.Timeout))
		}
		res, err := u.transport.RoundTrip(request)

		if err != nil {
			falcore.Error("Failed Ping to %v:%v: %v", u.Host, u.Port, err)
			return false, true
		} else {
			res.Body.Close()
		}
		if res.StatusCode == 200 {
			return true, true
		}
		falcore.Error("Failed Ping to %v:%v: %v", u.Host, u.Port, res.Status)
		// bad status
		return false, true
	}
	return false, false
}
<|endoftext|>


<commit_before>package stun

import (
	"encoding/binary"
	"fmt"

	"github.com/pkg/errors"
	"golang.org/x/net/ipv4"
)

// MessageClass of 0b00 is a request, a class of 0b01 is an
//   indication, a class of 0b10 is a success response, and a class of
//   0b11 is an error response.
// https://tools.ietf.org/html/rfc5389#section-6
type MessageClass byte

const (
	// ClassRequest describes a request method type
	ClassRequest MessageClass = 0x00
	// ClassIndication describes an indication method type
	ClassIndication MessageClass = 0x01
	// ClassSuccessResponse describes an success response method type
	ClassSuccessResponse MessageClass = 0x02
	// ClassErrorResponse describes an error response method type
	ClassErrorResponse MessageClass = 0x03
)

type Method uint16

const (
	MethodBinding          Method = 0x01 // STUN
	MethodSharedSecret     Method = 0x02 // STUN
	MethodAllocate         Method = 0x03 // TURN (Req/Rsp)
	MethodRefresh          Method = 0x04 // TURN (Req/Rsp)
	MethodSend             Method = 0x06 // TURN (Ind)
	MethodData             Method = 0x07 // TURN (Ind)
	MethodCreatePermission Method = 0x08 // TURN (Req/Rsp)
	MethodChannelBind      Method = 0x09 // TURN (Req/Rsp)
)

var messageClassName = map[MessageClass]string{
	ClassRequest:         "REQUEST",
	ClassIndication:      "INDICATION",
	ClassSuccessResponse: "SUCCESS-RESPONSE",
	ClassErrorResponse:   "ERROR-RESPONSE",
}

func (m MessageClass) String() string {
	s, err := messageClassName[m]
	if !err {
		// Falling back to hex representation.
		s = fmt.Sprintf("Unk 0x%x", uint16(m))
	}
	return s
}

var methodName = map[Method]string{
	MethodBinding:          "BINDING",
	MethodSharedSecret:     "SHARED-SECRET",
	MethodAllocate:         "ALLOCATE",
	MethodRefresh:          "REFRESH",
	MethodSend:             "SEND",
	MethodData:             "DATA",
	MethodCreatePermission: "CREATE-PERMISSION",
	MethodChannelBind:      "CHANNEL-BIND",
}

func (m Method) String() string {
	s, err := methodName[m]
	if !err {
		s = fmt.Sprintf("Unk 0x%x", uint16(m))
	}
	return s
}

//       0                   1                   2                   3
//       0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |0 0|     STUN Message Type     |         Message Length        |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |                         Magic Cookie                          |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |                                                               |
//      |                     Transaction ID (96 bits)                  |
//      |                                                               |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//

const (
	messageHeaderStart  int = 0
	messageHeaderLength int = 20
	messageLengthStart  int = 2
	messageLengthLength int = 2
	magicCookieStart    int = 4
	magicCookieLength   int = 4
	transactionIDStart  int = 4
	transactionIDLength int = 16
)

type Message struct {
	Class         MessageClass
	Method        Method
	Length        uint16
	TransactionID []byte
	Attributes    []*RawAttribute
	Raw           []byte
}

// The most significant 2 bits of every STUN message MUST be zeroes.
// This can be used to differentiate STUN packets from other protocols
// when STUN is multiplexed with other protocols on the same port.
// https://tools.ietf.org/html/rfc5389#section-6
func verifyStunHeaderMostSignificant2Bits(header []byte) bool {
	return (header[0] >> 6) == 0
}

func verifyMagicCookie(header []byte) error {
	const magicCookie = 0x2112A442
	c := header[magicCookieStart : magicCookieStart+magicCookieLength]
	if binary.BigEndian.Uint32(c) != magicCookie {
		return errors.Errorf("stun header magic cookie invalid; %v != %v (expected)", binary.BigEndian.Uint32(c), magicCookie)
	}
	return nil
}

// The message length MUST contain the size, in bytes, of the message
// not including the 20-byte STUN header.  Since all STUN attributes are
// padded to a multiple of 4 bytes, the last 2 bits of this field are
// always zero.  This provides another way to distinguish STUN packets
// from packets of other protocols.
// https://tools.ietf.org/html/rfc5389#section-6
func getMessageLength(header []byte) (uint16, error) {
	messageLength := binary.BigEndian.Uint16(header[messageLengthStart : messageLengthStart+messageLengthLength])
	if messageLength%4 != 0 {
		return 0, errors.Errorf("stun header message length must be a factor of 4 (%d)", messageLength)
	}

	return messageLength, nil
}

//  0                 1
//  2  3  4 5 6 7 8 9 0 1 2 3 4 5
//
// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
// |M |M |M|M|M|C|M|M|M|C|M|M|M|M|
// |11|10|9|8|7|1|6|5|4|0|3|2|1|0|
// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
const (
	c0Mask   = 0x10 // 0b10000
	c1Mask   = 0x01 // 0b00001
	c0ShiftR = 4    // R 0b10000 -> 0b00001
	c1ShiftL = 1    // L 0b00001 -> 0b00010

	m0Mask   = 0x0F // 0b00001111
	m4Mask   = 0xE0 // 0b11100000
	m7Mask   = 0x3E // 0b00111110
	m4ShiftR = 1    // R 0b01110000 -> 0b00111000
	m7ShiftL = 5    // L 0b00111110 -> 0b0000011111000000
)

func setMessageType(header []byte, class MessageClass, method Method) {
	m := uint16(method)
	c := uint16(class)

	mt := m & m0Mask
	// Make room for c0
	mt |= (m & (m4Mask >> m4ShiftR)) << 1
	mt |= (m & (m7Mask << 6)) << 2
	mt |= (c & 0x1) << 4
	mt |= (c >> 1) << 8

	binary.BigEndian.PutUint16(header[messageHeaderStart:], mt)
}

func getMessageType(header []byte) (MessageClass, Method) {
	mByte0 := header[0]
	mByte1 := header[1]

	c0 := (mByte1 & c0Mask) >> c0ShiftR
	c1 := (mByte0 & c1Mask) << c1ShiftL

	class := MessageClass(c1 | c0)

	m := (uint16(mByte0) & m7Mask) << m7ShiftL
	m |= uint16(mByte1 & m0Mask)
	m |= uint16((mByte1 & m4Mask) >> m4ShiftR)

	method := Method(m)

	return class, method
}

//  0                   1                   2                   3
//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |         Type                  |            Length             |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |                         Value (variable)                ....
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
func getAttribute(attribute []byte, offset int) *RawAttribute {
	typ := AttrType(binary.BigEndian.Uint16(attribute))
	len := binary.BigEndian.Uint16(attribute[attrLengthStart : attrLengthStart+attrLengthLength])
	pad := (attrLengthMultiple - (len % attrLengthMultiple)) % attrLengthMultiple
	return &RawAttribute{typ, len, attribute[attrValueStart : attrValueStart+len], pad, offset}
}

// TODO Break this apart, too big
func NewMessage(packet []byte) (*Message, error) {

	if len(packet) < 20 {
		return nil, errors.Errorf("stun header must be at least 20 bytes, was %d", len(packet))
	}

	header := packet[messageHeaderStart : messageHeaderStart+messageHeaderLength]

	if !verifyStunHeaderMostSignificant2Bits(header) {
		return nil, errors.New("stun header most significant 2 bits must equal 0b00")
	}

	err := verifyMagicCookie(header)
	if err != nil {
		return nil, errors.Wrap(err, "stun header invalid")
	}

	ml, err := getMessageLength(header)
	if err != nil {
		return nil, errors.Wrap(err, "stun header invalid")
	}

	if len(packet) != messageHeaderLength+int(ml) {
		return nil, errors.Errorf("stun header length invalid; %d != %d (expected)", messageHeaderLength+int(ml), len(packet))
	}

	t := header[transactionIDStart : transactionIDStart+transactionIDLength]

	class, method := getMessageType(header)

	ra := []*RawAttribute{}
	// TODO Check attr length <= attr slice remaining
	attr := packet[messageHeaderLength:]
	for len(attr) > 0 {
		a := getAttribute(attr, cap(packet)-cap(attr))
		attr = attr[attrValueStart+a.Length+a.Pad:]
		ra = append(ra, a)
	}

	m := Message{}
	m.Class = class
	m.Method = method
	m.Length = ml
	m.TransactionID = t[0:transactionIDLength]
	m.Attributes = ra
	m.Raw = packet

	return &m, nil
}

func (m *Message) GetOneAttribute(attrType AttrType) (*RawAttribute, bool) {
	for _, v := range m.Attributes {
		if v.Type == attrType {
			return v, true
		}
	}

	return nil, false
}

func (m *Message) GetAllAttributes(attrType AttrType) ([]*RawAttribute, bool) {
	var attrs []*RawAttribute
	for _, v := range m.Attributes {
		if v.Type == attrType {
			attrs = append(attrs, v)
		}
	}

	return attrs, len(attrs) > 0
}

func (m *Message) CommitLength() {
	binary.BigEndian.PutUint16(m.Raw[messageLengthStart:], uint16(m.Length))
}

func (m *Message) AddAttribute(attrType AttrType, v []byte) {

	ra := RawAttribute{
		Type:   attrType,
		Value:  v,
		Pad:    uint16(getAttrPadding(len(v))),
		Length: uint16(len(v)),
		Offset: int(m.Length),
	}

	a := make([]byte, attrHeaderLength+ra.Length+ra.Pad)

	binary.BigEndian.PutUint16(a, uint16(ra.Type))
	binary.BigEndian.PutUint16(a[attrLengthStart:attrLengthStart+attrLengthLength], ra.Length)

	copy(a[attrValueStart:], ra.Value)

	m.Attributes = append(m.Attributes, &ra)
	m.Raw = append(m.Raw, a...)
	m.Length += uint16(len(a))
	m.CommitLength()
}

func (m *Message) Pack() []byte {

	setMessageType(m.Raw[messageHeaderStart:2], m.Class, m.Method)
	m.CommitLength()
	copy(m.Raw[transactionIDStart:], m.TransactionID)

	return m.Raw
}

func BuildAndSend(conn *ipv4.PacketConn, addr *TransportAddr, class MessageClass, method Method, transactionID []byte, attrs ...Attribute) error {

	rsp, err := Build(class, method, transactionID, attrs...)
	if err != nil {
		return err
	}

	b := rsp.Pack()
	l, err := conn.WriteTo(b, nil, addr.Addr())
	if err != nil {
		return errors.Wrap(err, "failed writing to socket")
	}

	if l != len(b) {
		return errors.Errorf("packet write smaller than packet %d != %d (expected)", l, len(b))
	}

	return nil
}
<commit_msg>Add GenerateTransactionId<commit_after>package stun

import (
	"encoding/binary"
	"fmt"
	"math/rand"

	"github.com/pkg/errors"
	"golang.org/x/net/ipv4"
)

// MessageClass of 0b00 is a request, a class of 0b01 is an
//   indication, a class of 0b10 is a success response, and a class of
//   0b11 is an error response.
// https://tools.ietf.org/html/rfc5389#section-6
type MessageClass byte

const (
	// ClassRequest describes a request method type
	ClassRequest MessageClass = 0x00
	// ClassIndication describes an indication method type
	ClassIndication MessageClass = 0x01
	// ClassSuccessResponse describes an success response method type
	ClassSuccessResponse MessageClass = 0x02
	// ClassErrorResponse describes an error response method type
	ClassErrorResponse MessageClass = 0x03
)

type Method uint16

const (
	MethodBinding          Method = 0x01 // STUN
	MethodSharedSecret     Method = 0x02 // STUN
	MethodAllocate         Method = 0x03 // TURN (Req/Rsp)
	MethodRefresh          Method = 0x04 // TURN (Req/Rsp)
	MethodSend             Method = 0x06 // TURN (Ind)
	MethodData             Method = 0x07 // TURN (Ind)
	MethodCreatePermission Method = 0x08 // TURN (Req/Rsp)
	MethodChannelBind      Method = 0x09 // TURN (Req/Rsp)
)

var messageClassName = map[MessageClass]string{
	ClassRequest:         "REQUEST",
	ClassIndication:      "INDICATION",
	ClassSuccessResponse: "SUCCESS-RESPONSE",
	ClassErrorResponse:   "ERROR-RESPONSE",
}

func (m MessageClass) String() string {
	s, err := messageClassName[m]
	if !err {
		// Falling back to hex representation.
		s = fmt.Sprintf("Unk 0x%x", uint16(m))
	}
	return s
}

var methodName = map[Method]string{
	MethodBinding:          "BINDING",
	MethodSharedSecret:     "SHARED-SECRET",
	MethodAllocate:         "ALLOCATE",
	MethodRefresh:          "REFRESH",
	MethodSend:             "SEND",
	MethodData:             "DATA",
	MethodCreatePermission: "CREATE-PERMISSION",
	MethodChannelBind:      "CHANNEL-BIND",
}

func (m Method) String() string {
	s, err := methodName[m]
	if !err {
		s = fmt.Sprintf("Unk 0x%x", uint16(m))
	}
	return s
}

//       0                   1                   2                   3
//       0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |0 0|     STUN Message Type     |         Message Length        |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |                         Magic Cookie                          |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//      |                                                               |
//      |                     Transaction ID (96 bits)                  |
//      |                                                               |
//      +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//

const (
	messageHeaderStart  int = 0
	messageHeaderLength int = 20
	messageLengthStart  int = 2
	messageLengthLength int = 2
	magicCookieStart    int = 4
	magicCookieLength   int = 4
	transactionIDStart  int = 4
	transactionIDLength int = 16
)

type Message struct {
	Class         MessageClass
	Method        Method
	Length        uint16
	TransactionID []byte
	Attributes    []*RawAttribute
	Raw           []byte
}

// The most significant 2 bits of every STUN message MUST be zeroes.
// This can be used to differentiate STUN packets from other protocols
// when STUN is multiplexed with other protocols on the same port.
// https://tools.ietf.org/html/rfc5389#section-6
func verifyStunHeaderMostSignificant2Bits(header []byte) bool {
	return (header[0] >> 6) == 0
}

func verifyMagicCookie(header []byte) error {
	const magicCookie = 0x2112A442
	c := header[magicCookieStart : magicCookieStart+magicCookieLength]
	if binary.BigEndian.Uint32(c) != magicCookie {
		return errors.Errorf("stun header magic cookie invalid; %v != %v (expected)", binary.BigEndian.Uint32(c), magicCookie)
	}
	return nil
}

// The message length MUST contain the size, in bytes, of the message
// not including the 20-byte STUN header.  Since all STUN attributes are
// padded to a multiple of 4 bytes, the last 2 bits of this field are
// always zero.  This provides another way to distinguish STUN packets
// from packets of other protocols.
// https://tools.ietf.org/html/rfc5389#section-6
func getMessageLength(header []byte) (uint16, error) {
	messageLength := binary.BigEndian.Uint16(header[messageLengthStart : messageLengthStart+messageLengthLength])
	if messageLength%4 != 0 {
		return 0, errors.Errorf("stun header message length must be a factor of 4 (%d)", messageLength)
	}

	return messageLength, nil
}

//  0                 1
//  2  3  4 5 6 7 8 9 0 1 2 3 4 5
//
// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
// |M |M |M|M|M|C|M|M|M|C|M|M|M|M|
// |11|10|9|8|7|1|6|5|4|0|3|2|1|0|
// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
const (
	c0Mask   = 0x10 // 0b10000
	c1Mask   = 0x01 // 0b00001
	c0ShiftR = 4    // R 0b10000 -> 0b00001
	c1ShiftL = 1    // L 0b00001 -> 0b00010

	m0Mask   = 0x0F // 0b00001111
	m4Mask   = 0xE0 // 0b11100000
	m7Mask   = 0x3E // 0b00111110
	m4ShiftR = 1    // R 0b01110000 -> 0b00111000
	m7ShiftL = 5    // L 0b00111110 -> 0b0000011111000000
)

func setMessageType(header []byte, class MessageClass, method Method) {
	m := uint16(method)
	c := uint16(class)

	mt := m & m0Mask
	// Make room for c0
	mt |= (m & (m4Mask >> m4ShiftR)) << 1
	mt |= (m & (m7Mask << 6)) << 2
	mt |= (c & 0x1) << 4
	mt |= (c >> 1) << 8

	binary.BigEndian.PutUint16(header[messageHeaderStart:], mt)
}

func getMessageType(header []byte) (MessageClass, Method) {
	mByte0 := header[0]
	mByte1 := header[1]

	c0 := (mByte1 & c0Mask) >> c0ShiftR
	c1 := (mByte0 & c1Mask) << c1ShiftL

	class := MessageClass(c1 | c0)

	m := (uint16(mByte0) & m7Mask) << m7ShiftL
	m |= uint16(mByte1 & m0Mask)
	m |= uint16((mByte1 & m4Mask) >> m4ShiftR)

	method := Method(m)

	return class, method
}

//  0                   1                   2                   3
//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |         Type                  |            Length             |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |                         Value (variable)                ....
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
func getAttribute(attribute []byte, offset int) *RawAttribute {
	typ := AttrType(binary.BigEndian.Uint16(attribute))
	len := binary.BigEndian.Uint16(attribute[attrLengthStart : attrLengthStart+attrLengthLength])
	pad := (attrLengthMultiple - (len % attrLengthMultiple)) % attrLengthMultiple
	return &RawAttribute{typ, len, attribute[attrValueStart : attrValueStart+len], pad, offset}
}

// TODO Break this apart, too big
func NewMessage(packet []byte) (*Message, error) {

	if len(packet) < 20 {
		return nil, errors.Errorf("stun header must be at least 20 bytes, was %d", len(packet))
	}

	header := packet[messageHeaderStart : messageHeaderStart+messageHeaderLength]

	if !verifyStunHeaderMostSignificant2Bits(header) {
		return nil, errors.New("stun header most significant 2 bits must equal 0b00")
	}

	err := verifyMagicCookie(header)
	if err != nil {
		return nil, errors.Wrap(err, "stun header invalid")
	}

	ml, err := getMessageLength(header)
	if err != nil {
		return nil, errors.Wrap(err, "stun header invalid")
	}

	if len(packet) != messageHeaderLength+int(ml) {
		return nil, errors.Errorf("stun header length invalid; %d != %d (expected)", messageHeaderLength+int(ml), len(packet))
	}

	t := header[transactionIDStart : transactionIDStart+transactionIDLength]

	class, method := getMessageType(header)

	ra := []*RawAttribute{}
	// TODO Check attr length <= attr slice remaining
	attr := packet[messageHeaderLength:]
	for len(attr) > 0 {
		a := getAttribute(attr, cap(packet)-cap(attr))
		attr = attr[attrValueStart+a.Length+a.Pad:]
		ra = append(ra, a)
	}

	m := Message{}
	m.Class = class
	m.Method = method
	m.Length = ml
	m.TransactionID = t[0:transactionIDLength]
	m.Attributes = ra
	m.Raw = packet

	return &m, nil
}

func (m *Message) GetOneAttribute(attrType AttrType) (*RawAttribute, bool) {
	for _, v := range m.Attributes {
		if v.Type == attrType {
			return v, true
		}
	}

	return nil, false
}

func (m *Message) GetAllAttributes(attrType AttrType) ([]*RawAttribute, bool) {
	var attrs []*RawAttribute
	for _, v := range m.Attributes {
		if v.Type == attrType {
			attrs = append(attrs, v)
		}
	}

	return attrs, len(attrs) > 0
}

func (m *Message) CommitLength() {
	binary.BigEndian.PutUint16(m.Raw[messageLengthStart:], uint16(m.Length))
}

func (m *Message) AddAttribute(attrType AttrType, v []byte) {

	ra := RawAttribute{
		Type:   attrType,
		Value:  v,
		Pad:    uint16(getAttrPadding(len(v))),
		Length: uint16(len(v)),
		Offset: int(m.Length),
	}

	a := make([]byte, attrHeaderLength+ra.Length+ra.Pad)

	binary.BigEndian.PutUint16(a, uint16(ra.Type))
	binary.BigEndian.PutUint16(a[attrLengthStart:attrLengthStart+attrLengthLength], ra.Length)

	copy(a[attrValueStart:], ra.Value)

	m.Attributes = append(m.Attributes, &ra)
	m.Raw = append(m.Raw, a...)
	m.Length += uint16(len(a))
	m.CommitLength()
}

func (m *Message) Pack() []byte {

	setMessageType(m.Raw[messageHeaderStart:2], m.Class, m.Method)
	m.CommitLength()
	copy(m.Raw[transactionIDStart:], m.TransactionID)

	return m.Raw
}

func BuildAndSend(conn *ipv4.PacketConn, addr *TransportAddr, class MessageClass, method Method, transactionID []byte, attrs ...Attribute) error {

	rsp, err := Build(class, method, transactionID, attrs...)
	if err != nil {
		return err
	}

	b := rsp.Pack()
	l, err := conn.WriteTo(b, nil, addr.Addr())
	if err != nil {
		return errors.Wrap(err, "failed writing to socket")
	}

	if l != len(b) {
		return errors.Errorf("packet write smaller than packet %d != %d (expected)", l, len(b))
	}

	return nil
}

func GenerateTransactionId() []byte {
	randSeq := func(n int) string {
		letters := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
		b := make([]rune, n)
		for i := range b {
			b[i] = letters[rand.Intn(len(letters))]
		}
		return string(b)
	}

	transactionID := []byte(randSeq(16))
	transactionID[0] = 33
	transactionID[1] = 18
	transactionID[2] = 164
	transactionID[3] = 66
	return transactionID
}
<|endoftext|>


<commit_before>package f2k

import (
  "fmt"
  "net/http"
  "math/rand"
  "time"
  "strconv"
  "encoding/json"
  "log"

  "appengine"
  "appengine/urlfetch"
)

type Results struct {
  Result []Story `json:"hits"`
}

type Story struct {
  Title, Url string
}

func init() {
  rand.Seed( time.Now().UTC().UnixNano())
  http.HandleFunc("/", handler)
  http.HandleFunc("/pod", pod_handler)

  http.ListenAndServe(":8080", nil)
}

func handler(w http.ResponseWriter, r *http.Request) {
  fmt.Fprint(w, "f2k is an api that provides random HN posts when requested. Fork here - https://github.com/avinoth/f2k")
}

func pod_handler(w http.ResponseWriter, r *http.Request) {
  w.Header().Set("Content-Type", "application/json")

  var results Results

  hn_url := "http://hn.algolia.com/api/v1/"
  page := rand.Intn(50)
  points := randomInt(200, 600)
  url := hn_url + "search?tags=story&numericFilters=points>" + strconv.Itoa(points) + "&page=" + strconv.Itoa(page)

  resp, err := makeRequest(url, r)

  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  err = json.NewDecoder(resp.Body).Decode(&results)

  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  res, err := json.Marshal(results.Result[0])
  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  fmt.Fprint(w, string(res))
}

func randomInt(min, max int) int {
  return min + rand.Intn(max - min)
}

func makeRequest(url string, r *http.Request) (*http.Response, error) {
  c := appengine.NewContext(r)
  client := urlfetch.Client(c)

  response, err := client.Get(url)
  if err != nil {
    return nil, err
  }

  if response.StatusCode == http.StatusNotFound {
    return nil, fmt.Errorf(http.StatusText(http.StatusNotFound))
  }
  return response, nil
}
<commit_msg>randomize index<commit_after>package f2k

import (
  "fmt"
  "net/http"
  "math/rand"
  "time"
  "strconv"
  "encoding/json"
  "log"

  "appengine"
  "appengine/urlfetch"
)

type Results struct {
  Result []Story `json:"hits"`
}

type Story struct {
  Title, Url string
}

func init() {
  rand.Seed( time.Now().UTC().UnixNano())
  http.HandleFunc("/", handler)
  http.HandleFunc("/pod", pod_handler)

  http.ListenAndServe(":8080", nil)
}

func handler(w http.ResponseWriter, r *http.Request) {
  fmt.Fprint(w, "f2k is an api that provides random HN posts when requested. Fork here - https://github.com/avinoth/f2k")
}

func pod_handler(w http.ResponseWriter, r *http.Request) {
  w.Header().Set("Content-Type", "application/json")

  var results Results

  hn_url := "http://hn.algolia.com/api/v1/"
  page := rand.Intn(50)
  points := randomInt(200, 600)
  url := hn_url + "search?tags=story&numericFilters=points>" + strconv.Itoa(points) + "&page=" + strconv.Itoa(page)

  resp, err := makeRequest(url, r)

  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  err = json.NewDecoder(resp.Body).Decode(&results)

  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  idx := randomInt(0, 19)
  res, err := json.Marshal(results.Result[idx])
  if err != nil {
    fmt.Fprint(w, "Something Went Wrong!")
    log.Fatal(err)
  }

  fmt.Fprint(w, string(res))
}

func randomInt(min, max int) int {
  return min + rand.Intn(max - min)
}

func makeRequest(url string, r *http.Request) (*http.Response, error) {
  c := appengine.NewContext(r)
  client := urlfetch.Client(c)

  response, err := client.Get(url)
  if err != nil {
    return nil, err
  }

  if response.StatusCode == http.StatusNotFound {
    return nil, fmt.Errorf(http.StatusText(http.StatusNotFound))
  }
  return response, nil
}
<|endoftext|>


<commit_before>package fakku

import (
	"encoding/json"
	"fmt"
	"strconv"
)

const (
	ErrorContentDoesntExist = "Content doesn't exist"
	ErrorUnknownJsonData    = "Got unknown json data back from content request. API Change?"
	ErrorUnknownJsonLayout  = "Got an unknown layout back from content request. API Change?"
)

type Content struct {
	Name         string
	Url          string
	Description  string
	Language     string
	Category     string
	Date         float64
	FileSize     float64
	Favorites    float64
	CommentCount float64
	Pages        float64
	Poster       string
	PosterUrl    string
	Tags         []*Attribute `json:"content_tags"`
	Translators  []*Attribute `json:"content_translators"`
	Series       []*Attribute `json:"content_series"`
	Artists      []*Attribute `json:"content_artists"`
	Images       struct {
		Cover  string
		Sample string
	}
}
type ContentList []Content

func (c *Content) UnmarshalJSON(b []byte) error {
	var f interface{}
	json.Unmarshal(b, &f)
	switch f.(type) {
	case map[string]interface{}:
		m := f.(map[string]interface{})
		contents := m["content"]
		switch contents.(type) {
		case map[string]interface{}:
			v := contents.(map[string]interface{})
			c.populateContent(v)
			return nil
		case []interface{}:
			q := contents.([]interface{})
			if len(q) == 0 {
				// doesn't exist
				return fmt.Errorf(ErrorContentDoesntExist)
			} else {
				return fmt.Errorf(ErrorUnknownJsonData)
			}
		default:
			return fmt.Errorf(ErrorUnknownJsonLayout)
		}
	case []interface{}:
		q := f.([]interface{})
		if len(q) == 0 {
			// doesn't exist
			return fmt.Errorf(ErrorContentDoesntExist)
		} else {
			return fmt.Errorf(ErrorUnknownJsonData)
		}
	default:
		return fmt.Errorf(ErrorUnknownJsonLayout)
	}
}

func (c *Content) populateContent(v map[string]interface{}) {
	c.Name = v["content_name"].(string)
	c.Url = v["content_url"].(string)
	c.Description = v["content_description"].(string)
	c.Language = v["content_language"].(string)
	c.Category = v["content_category"].(string)
	c.Date = v["content_date"].(float64)
	c.FileSize = v["content_filesize"].(float64)
	c.Favorites = v["content_favorites"].(float64)
	c.CommentCount = v["content_comments"].(float64)
	c.Pages = v["content_pages"].(float64)
	c.Poster = v["content_poster"].(string)
	c.PosterUrl = v["content_poster_url"].(string)
	c.Tags = constructAttributeFields(v, "content_tags")
	c.Translators = constructAttributeFields(v, "content_translators")
	c.Series = constructAttributeFields(v, "content_series")
	c.Artists = constructAttributeFields(v, "content_artists")

	tmp := v["content_images"]
	z := tmp.(map[string]interface{})
	c.Images.Cover = z["cover"].(string)
	c.Images.Sample = z["sample"].(string)
}

func GetContent(category, name string) (*Content, error) {
	var c Content
	url := contentApiFunction{Category: category, Name: name}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

type contentApiFunction struct {
	Category string
	Name     string
}

func (a contentApiFunction) Construct() string {
	return fmt.Sprintf("%s/%s/%s", ApiHeader, a.Category, a.Name)
}

type contentCommentApiFunction struct {
	contentApiFunction
	TopComments bool
	SupportsPagination
}

func (a contentCommentApiFunction) Construct() string {
	base := fmt.Sprintf("%s/comments", a.contentApiFunction.Construct())
	if a.TopComments {
		return fmt.Sprintf("%s/top", base)
	} else {
		return PaginateString(base, a.Page)
	}
}

func getContentCommentsGeneric(url ApiFunction) (*Comments, error) {
	var c Comments
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}
func ContentComments(category, name string) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	return getContentCommentsGeneric(url)
}

func ContentCommentsPage(category, name string, page uint) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		SupportsPagination: SupportsPagination{Page: page},
	}
	return getContentCommentsGeneric(url)
}

func ContentTopComments(category, name string) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		TopComments: true,
	}
	return getContentCommentsGeneric(url)
}

func (this *Content) TopComments() (*Comments, error) {
	return ContentTopComments(this.Category, this.Name)
}
func (this *Content) CommentsPage(page uint) (*Comments, error) {
	return ContentCommentsPage(this.Category, this.Name, page)
}
func (this *Content) Comments() (*Comments, error) {
	return ContentComments(this.Category, this.Name)
}

type Attribute struct {
	Attribute     string `json:"attribute"`
	AttributeLink string `json:"attribute_link"`
}

func NewAttribute(c map[string]interface{}) *Attribute {
	return &Attribute{
		Attribute:     c["attribute"].(string),
		AttributeLink: c["attribute_link"].(string),
	}
}

func (a *Attribute) String() string {
	return a.Attribute
}

func constructAttributeFields(c map[string]interface{}, field string) []*Attribute {
	try, ok := c[field]
	if !ok {
		return nil
	}
	tmp := try.([]interface{})
	size := len(tmp)
	attrs := make([]*Attribute, size)
	for i := 0; i < size; i++ {
		attrs[i] = NewAttribute(tmp[i].(map[string]interface{}))
	}
	return attrs
}

type Comment struct {
	Id         float64 `json:"comment_id"`
	AttachedId string  `json:"comment_attached_id"`
	Poster     string  `json:"comment_poster"`
	PosterUrl  string  `json:"comment_poster_url"`
	Reputation float64 `json:"comment_reputation"`
	Text       string  `json:"comment_text"`
	Date       float64 `json:"comment_date"`
}

type Comments struct {
	Comments   []*Comment `json:"comments"`
	PageNumber float64    `json:"page"`
	Total      float64    `json:"total"`
	Pages      float64    `json:"pages"`
}

type ReadOnlineContent struct {
	Content *Content `json:"content"`
	Pages   []*Page  `json:"pages"`
}

func (r *ReadOnlineContent) UnmarshalJSON(b []byte) error {
	var f interface{}
	if err := json.Unmarshal(b, &r.Content); err != nil {
		return err
	}
	json.Unmarshal(b, &f)
	// need to check and make sure that the content exists
	switch f.(type) {
	case map[string]interface{}:
		m := f.(map[string]interface{})
		pages := m["pages"]
		v := pages.(map[string]interface{})
		r.Pages = make([]*Page, len(v))
		for i := 0; i < len(v); i++ {
			ind := strconv.Itoa(i + 1)
			r.Pages[i] = NewPage(ind, v[ind].(map[string]interface{}))
		}
		return nil
	case []interface{}:
		q := f.([]interface{})
		if len(q) == 0 {
			// doesn't exist
			return fmt.Errorf("Content doesn't exist")
		} else {
			return fmt.Errorf("Got unknown json data back from content request. API Change?")
		}
	default:
		return fmt.Errorf("Got an unknown layout back from content request. API Change?")
	}
}

type Page struct {
	Id    string
	Thumb string
	Image string
}

func NewPage(id string, c map[string]interface{}) *Page {
	return &Page{
		Id:    id,
		Thumb: c["thumb"].(string),
		Image: c["image"].(string),
	}
}

type ContentReadOnlineApiFunction struct {
	contentApiFunction
}

func (a ContentReadOnlineApiFunction) Construct() string {
	return fmt.Sprintf("%s/read", a.contentApiFunction.Construct())
}

func GetContentReadOnline(category, name string) (*ReadOnlineContent, error) {
	var c ReadOnlineContent
	url := ContentReadOnlineApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

func GetContentDownloads(category, name string) (*DownloadContent, error) {
	var c DownloadContent
	url := ContentDownloadsApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

type ContentDownloadsApiFunction struct {
	contentApiFunction
}

func (a ContentDownloadsApiFunction) Construct() string {
	return fmt.Sprintf("%s/download", a.contentApiFunction.Construct())
}

type DownloadContent struct {
	Downloads []*Download `json:"downloads"`
	Total     uint        `json:"total"`
}

func (this *DownloadContent) HasDownloads() bool {
	return this.Total > 0
}

type Download struct {
	Type          string  `json:"download_type"`
	Url           string  `json:"download_url"`
	Info          string  `json:"download_info"`
	DownloadCount float64 `json:"download_count"`
	RawTime       float64 `json:"download_time"`
	Poster        string  `json:"download_poster"`
	RawPosterUrl  string  `json:"download_poster_url"`
}

type contentRelatedApiFunction struct {
	contentApiFunction
	SupportsPagination
}

func (a contentRelatedApiFunction) Construct() string {
	base := fmt.Sprintf("%s/related", a.contentApiFunction.Construct())
	return PaginateString(base, a.Page)
}

type RelatedContent struct {
	Related ContentList `json:"related"`
	Total   uint        `json:"total"`
	Pages   uint        `json:"pages"`
}

func GetRelatedContentAll(category, name string) (*RelatedContent, error) {
	return GetRelatedContent(category, name, 0)
}

func GetRelatedContent(category, name string, page uint) (*RelatedContent, error) {
	var c RelatedContent
	url := contentRelatedApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		SupportsPagination: SupportsPagination{Page: page},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}
func (this *Content) RelatedContent() (*RelatedContent, error) {
	return GetRelatedContentAll(this.Category, this.Name)
}
func (this *Content) RelatedContentPage(page uint) (*RelatedContent, error) {
	return GetRelatedContent(this.Category, this.Name, page)
}

func (c *RelatedContent) UnmarshalJSON(b []byte) error {
	// slightly different
	var f interface{}
	json.Unmarshal(b, &f)
	m := f.(map[string]interface{})
	related := m["related"]
	v := related.([]interface{})
	c.Related = make(ContentList, len(v))
	for i := 0; i < len(v); i++ {
		c.Related[i].populateContent(v[i].(map[string]interface{}))
	}
	c.Total = uint(m["total"].(float64))
	c.Pages = uint(m["pages"].(float64))
	return nil
}
<commit_msg>Added methods to the Download type<commit_after>package fakku

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"time"
)

const (
	ErrorContentDoesntExist = "Content doesn't exist"
	ErrorUnknownJsonData    = "Got unknown json data back from content request. API Change?"
	ErrorUnknownJsonLayout  = "Got an unknown layout back from content request. API Change?"
)

type Content struct {
	Name         string
	Url          string
	Description  string
	Language     string
	Category     string
	Date         float64
	FileSize     float64
	Favorites    float64
	CommentCount float64
	Pages        float64
	Poster       string
	PosterUrl    string
	Tags         []*Attribute `json:"content_tags"`
	Translators  []*Attribute `json:"content_translators"`
	Series       []*Attribute `json:"content_series"`
	Artists      []*Attribute `json:"content_artists"`
	Images       struct {
		Cover  string
		Sample string
	}
}
type ContentList []Content

func (c *Content) UnmarshalJSON(b []byte) error {
	var f interface{}
	json.Unmarshal(b, &f)
	switch f.(type) {
	case map[string]interface{}:
		m := f.(map[string]interface{})
		contents := m["content"]
		switch contents.(type) {
		case map[string]interface{}:
			v := contents.(map[string]interface{})
			c.populateContent(v)
			return nil
		case []interface{}:
			q := contents.([]interface{})
			if len(q) == 0 {
				// doesn't exist
				return fmt.Errorf(ErrorContentDoesntExist)
			} else {
				return fmt.Errorf(ErrorUnknownJsonData)
			}
		default:
			return fmt.Errorf(ErrorUnknownJsonLayout)
		}
	case []interface{}:
		q := f.([]interface{})
		if len(q) == 0 {
			// doesn't exist
			return fmt.Errorf(ErrorContentDoesntExist)
		} else {
			return fmt.Errorf(ErrorUnknownJsonData)
		}
	default:
		return fmt.Errorf(ErrorUnknownJsonLayout)
	}
}

func (c *Content) populateContent(v map[string]interface{}) {
	c.Name = v["content_name"].(string)
	c.Url = v["content_url"].(string)
	c.Description = v["content_description"].(string)
	c.Language = v["content_language"].(string)
	c.Category = v["content_category"].(string)
	c.Date = v["content_date"].(float64)
	c.FileSize = v["content_filesize"].(float64)
	c.Favorites = v["content_favorites"].(float64)
	c.CommentCount = v["content_comments"].(float64)
	c.Pages = v["content_pages"].(float64)
	c.Poster = v["content_poster"].(string)
	c.PosterUrl = v["content_poster_url"].(string)
	c.Tags = constructAttributeFields(v, "content_tags")
	c.Translators = constructAttributeFields(v, "content_translators")
	c.Series = constructAttributeFields(v, "content_series")
	c.Artists = constructAttributeFields(v, "content_artists")

	tmp := v["content_images"]
	z := tmp.(map[string]interface{})
	c.Images.Cover = z["cover"].(string)
	c.Images.Sample = z["sample"].(string)
}

func GetContent(category, name string) (*Content, error) {
	var c Content
	url := contentApiFunction{Category: category, Name: name}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

type contentApiFunction struct {
	Category string
	Name     string
}

func (a contentApiFunction) Construct() string {
	return fmt.Sprintf("%s/%s/%s", ApiHeader, a.Category, a.Name)
}

type contentCommentApiFunction struct {
	contentApiFunction
	TopComments bool
	SupportsPagination
}

func (a contentCommentApiFunction) Construct() string {
	base := fmt.Sprintf("%s/comments", a.contentApiFunction.Construct())
	if a.TopComments {
		return fmt.Sprintf("%s/top", base)
	} else {
		return PaginateString(base, a.Page)
	}
}

func getContentCommentsGeneric(url ApiFunction) (*Comments, error) {
	var c Comments
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}
func ContentComments(category, name string) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	return getContentCommentsGeneric(url)
}

func ContentCommentsPage(category, name string, page uint) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		SupportsPagination: SupportsPagination{Page: page},
	}
	return getContentCommentsGeneric(url)
}

func ContentTopComments(category, name string) (*Comments, error) {
	url := contentCommentApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		TopComments: true,
	}
	return getContentCommentsGeneric(url)
}

func (this *Content) TopComments() (*Comments, error) {
	return ContentTopComments(this.Category, this.Name)
}
func (this *Content) CommentsPage(page uint) (*Comments, error) {
	return ContentCommentsPage(this.Category, this.Name, page)
}
func (this *Content) Comments() (*Comments, error) {
	return ContentComments(this.Category, this.Name)
}

type Attribute struct {
	Attribute     string `json:"attribute"`
	AttributeLink string `json:"attribute_link"`
}

func NewAttribute(c map[string]interface{}) *Attribute {
	return &Attribute{
		Attribute:     c["attribute"].(string),
		AttributeLink: c["attribute_link"].(string),
	}
}

func (a *Attribute) String() string {
	return a.Attribute
}

func constructAttributeFields(c map[string]interface{}, field string) []*Attribute {
	try, ok := c[field]
	if !ok {
		return nil
	}
	tmp := try.([]interface{})
	size := len(tmp)
	attrs := make([]*Attribute, size)
	for i := 0; i < size; i++ {
		attrs[i] = NewAttribute(tmp[i].(map[string]interface{}))
	}
	return attrs
}

type Comment struct {
	Id         float64 `json:"comment_id"`
	AttachedId string  `json:"comment_attached_id"`
	Poster     string  `json:"comment_poster"`
	PosterUrl  string  `json:"comment_poster_url"`
	Reputation float64 `json:"comment_reputation"`
	Text       string  `json:"comment_text"`
	Date       float64 `json:"comment_date"`
}

type Comments struct {
	Comments   []*Comment `json:"comments"`
	PageNumber float64    `json:"page"`
	Total      float64    `json:"total"`
	Pages      float64    `json:"pages"`
}

type ReadOnlineContent struct {
	Content *Content `json:"content"`
	Pages   []*Page  `json:"pages"`
}

func (r *ReadOnlineContent) UnmarshalJSON(b []byte) error {
	var f interface{}
	if err := json.Unmarshal(b, &r.Content); err != nil {
		return err
	}
	json.Unmarshal(b, &f)
	// need to check and make sure that the content exists
	switch f.(type) {
	case map[string]interface{}:
		m := f.(map[string]interface{})
		pages := m["pages"]
		v := pages.(map[string]interface{})
		r.Pages = make([]*Page, len(v))
		for i := 0; i < len(v); i++ {
			ind := strconv.Itoa(i + 1)
			r.Pages[i] = NewPage(ind, v[ind].(map[string]interface{}))
		}
		return nil
	case []interface{}:
		q := f.([]interface{})
		if len(q) == 0 {
			// doesn't exist
			return fmt.Errorf("Content doesn't exist")
		} else {
			return fmt.Errorf("Got unknown json data back from content request. API Change?")
		}
	default:
		return fmt.Errorf("Got an unknown layout back from content request. API Change?")
	}
}

type Page struct {
	Id    string
	Thumb string
	Image string
}

func NewPage(id string, c map[string]interface{}) *Page {
	return &Page{
		Id:    id,
		Thumb: c["thumb"].(string),
		Image: c["image"].(string),
	}
}

type ContentReadOnlineApiFunction struct {
	contentApiFunction
}

func (a ContentReadOnlineApiFunction) Construct() string {
	return fmt.Sprintf("%s/read", a.contentApiFunction.Construct())
}

func GetContentReadOnline(category, name string) (*ReadOnlineContent, error) {
	var c ReadOnlineContent
	url := ContentReadOnlineApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

func GetContentDownloads(category, name string) (*DownloadContent, error) {
	var c DownloadContent
	url := ContentDownloadsApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}

type ContentDownloadsApiFunction struct {
	contentApiFunction
}

func (a ContentDownloadsApiFunction) Construct() string {
	return fmt.Sprintf("%s/download", a.contentApiFunction.Construct())
}

type DownloadContent struct {
	Downloads []*Download `json:"downloads"`
	Total     uint        `json:"total"`
}

func (this *DownloadContent) HasDownloads() bool {
	return this.Total > 0
}

type Download struct {
	Type          string  `json:"download_type"`
	RawUrl        string  `json:"download_url"`
	Info          string  `json:"download_info"`
	DownloadCount float64 `json:"download_count"`
	RawTime       float64 `json:"download_time"`
	Poster        string  `json:"download_poster"`
	RawPosterUrl  string  `json:"download_poster_url"`
}

func (this *Download) Url() (*url.URL, error) {
	return url.Parse(this.RawUrl)
}
func (this *Download) PosterUrl() (*url.URL, error) {
	return url.Parse(this.RawPosterUrl)
}
func (this *Download) Time() time.Time {
	return time.Unix(int64(this.RawTime), 0)
}

type contentRelatedApiFunction struct {
	contentApiFunction
	SupportsPagination
}

func (a contentRelatedApiFunction) Construct() string {
	base := fmt.Sprintf("%s/related", a.contentApiFunction.Construct())
	return PaginateString(base, a.Page)
}

type RelatedContent struct {
	Related ContentList `json:"related"`
	Total   uint        `json:"total"`
	Pages   uint        `json:"pages"`
}

func GetRelatedContentAll(category, name string) (*RelatedContent, error) {
	return GetRelatedContent(category, name, 0)
}

func GetRelatedContent(category, name string, page uint) (*RelatedContent, error) {
	var c RelatedContent
	url := contentRelatedApiFunction{
		contentApiFunction: contentApiFunction{
			Category: category,
			Name:     name,
		},
		SupportsPagination: SupportsPagination{Page: page},
	}
	if err := ApiCall(url, &c); err != nil {
		return nil, err
	} else {
		return &c, nil
	}
}
func (this *Content) RelatedContent() (*RelatedContent, error) {
	return GetRelatedContentAll(this.Category, this.Name)
}
func (this *Content) RelatedContentPage(page uint) (*RelatedContent, error) {
	return GetRelatedContent(this.Category, this.Name, page)
}

func (c *RelatedContent) UnmarshalJSON(b []byte) error {
	// slightly different
	var f interface{}
	json.Unmarshal(b, &f)
	m := f.(map[string]interface{})
	related := m["related"]
	v := related.([]interface{})
	c.Related = make(ContentList, len(v))
	for i := 0; i < len(v); i++ {
		c.Related[i].populateContent(v[i].(map[string]interface{}))
	}
	c.Total = uint(m["total"].(float64))
	c.Pages = uint(m["pages"].(float64))
	return nil
}
<|endoftext|>


<commit_before>package docker

import (
	"fmt"
	"path/filepath"

	"github.com/arschles/gci/config"
	"github.com/arschles/gci/util/docker/build"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/pborman/uuid"
)

const (
	containerOutDir = "/gobin"
)

func goxOutputTpl(binPath string) string {
	return fmt.Sprintf("%s_{{.OS}}_{{.Arch}}", binPath)
}

func imageName(crossCompile bool) string {
	if crossCompile {
		return GoxImage
	}
	return GolangImage
}

func command(crossCompile bool, binaryPath string) []string {
	if crossCompile {
		return []string{"gox", "-output", goxOutputTpl(binaryPath)}
	}
	return []string{"go", "build", "-o", binaryPath}
}

// Build runs the build of rootDir inside a Docker container, putting binaries into outDir
func Build(
	dockerCl *docker.Client,
	rootDir,
	outDir,
	packageName,
	containerGoPath string,
	cfg *config.File,
	logsCh chan<- build.Log,
	resultCh chan<- int,
	errCh chan<- error) {

	projName := filepath.Base(rootDir)
	imgName := imageName(cfg.Build.CrossCompile)
	containerName := fmt.Sprintf("gci-build-%s-%s", projName, uuid.New())
	logsCh <- build.LogFromString("Creating container %s to build %s", containerName, packageName)

	binaryName := cfg.Build.GetOutputBinary(projName)
	cmd := command(cfg.Build.CrossCompile, fmt.Sprintf("%s/%s", containerOutDir, binaryName))
	env := cfg.Build.Env

	containerWorkDir := fmt.Sprintf("%s/src/%s", containerGoPath, packageName)

	mounts := []docker.Mount{
		{
			Name:        "source_dir",
			Source:      rootDir,
			Destination: containerWorkDir,
			Mode:        "r",
		},
		{
			Name:        "dest_dir",
			Source:      outDir,
			Destination: containerOutDir,
			Mode:        "w",
		},
	}
	createContainerOpts, hostConfig := CreateAndStartContainerOpts(
		imgName,
		containerName,
		cmd,
		env,
		mounts,
		containerGoPath,
		containerWorkDir,
	)
	container, err := dockerCl.CreateContainer(createContainerOpts)
	if err != nil {
		errCh <- fmt.Errorf("error creating container (%s)", err)
		return
	}

	logsCh <- build.LogFromString(CmdStr(createContainerOpts, hostConfig))

	if err := dockerCl.StartContainer(container.ID, &hostConfig); err != nil {
		errCh <- fmt.Errorf("error starting container (%s)", err)
		return
	}

	stdOut := build.NewChanWriter(logsCh)
	stdErr := build.NewChanWriter(logsCh)
	attachOpts := AttachToContainerOpts(container.ID, stdOut, stdErr)
	waitCodeCh, waitErrCh, err := AttachAndWait(dockerCl, container.ID, attachOpts)

	if err != nil {
		errCh <- fmt.Errorf("error attaching to the build container (%s)", err)
		return
	}

	select {
	case err := <-waitErrCh:
		errCh <- fmt.Errorf("error waiting for the build container to finish (%s)", err)
		return
	case code := <-waitCodeCh:
		resultCh <- code
	}
}
<commit_msg>remove the build container after use<commit_after>package docker

import (
	"fmt"
	"path/filepath"

	"github.com/arschles/gci/config"
	"github.com/arschles/gci/util/docker/build"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/pborman/uuid"
)

const (
	containerOutDir = "/gobin"
)

func goxOutputTpl(binPath string) string {
	return fmt.Sprintf("%s_{{.OS}}_{{.Arch}}", binPath)
}

func imageName(crossCompile bool) string {
	if crossCompile {
		return GoxImage
	}
	return GolangImage
}

func command(crossCompile bool, binaryPath string) []string {
	if crossCompile {
		return []string{"gox", "-output", goxOutputTpl(binaryPath)}
	}
	return []string{"go", "build", "-o", binaryPath}
}

// Build runs the build of rootDir inside a Docker container, putting binaries into outDir
func Build(
	dockerCl *docker.Client,
	rootDir,
	outDir,
	packageName,
	containerGoPath string,
	cfg *config.File,
	logsCh chan<- build.Log,
	resultCh chan<- int,
	errCh chan<- error) {

	projName := filepath.Base(rootDir)
	imgName := imageName(cfg.Build.CrossCompile)
	containerName := fmt.Sprintf("gci-build-%s-%s", projName, uuid.New())
	logsCh <- build.LogFromString("Creating container %s to build %s", containerName, packageName)

	binaryName := cfg.Build.GetOutputBinary(projName)
	cmd := command(cfg.Build.CrossCompile, fmt.Sprintf("%s/%s", containerOutDir, binaryName))
	env := cfg.Build.Env

	containerWorkDir := fmt.Sprintf("%s/src/%s", containerGoPath, packageName)

	mounts := []docker.Mount{
		{
			Name:        "source_dir",
			Source:      rootDir,
			Destination: containerWorkDir,
			Mode:        "r",
		},
		{
			Name:        "dest_dir",
			Source:      outDir,
			Destination: containerOutDir,
			Mode:        "w",
		},
	}
	createContainerOpts, hostConfig := CreateAndStartContainerOpts(
		imgName,
		containerName,
		cmd,
		env,
		mounts,
		containerGoPath,
		containerWorkDir,
	)
	container, err := dockerCl.CreateContainer(createContainerOpts)
	if err != nil {
		errCh <- fmt.Errorf("error creating container (%s)", err)
		return
	}

	logsCh <- build.LogFromString(CmdStr(createContainerOpts, hostConfig))

	if err := dockerCl.StartContainer(container.ID, &hostConfig); err != nil {
		errCh <- fmt.Errorf("error starting container (%s)", err)
		return
	}

	defer func() {
		if err := dockerCl.RemoveContainer(docker.RemoveContainerOpts{ID: container.ID, Force: true}); err != nil {
			log.Printf("Error removing build container %s (%s)", container.ID, err)
		}
	}()

	stdOut := build.NewChanWriter(logsCh)
	stdErr := build.NewChanWriter(logsCh)
	attachOpts := AttachToContainerOpts(container.ID, stdOut, stdErr)
	waitCodeCh, waitErrCh, err := AttachAndWait(dockerCl, container.ID, attachOpts)

	if err != nil {
		errCh <- fmt.Errorf("error attaching to the build container (%s)", err)
		return
	}

	select {
	case err := <-waitErrCh:
		errCh <- fmt.Errorf("error waiting for the build container to finish (%s)", err)
		return
	case code := <-waitCodeCh:
		resultCh <- code
	}
}
<|endoftext|>