<commit_before>// Copyright 2013 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package protorpc

import (
	"fmt"
	"hash/crc32"
	"io"

	wire "github.com/chai2010/protorpc/wire.pb"
	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
)

func maxUint32(a, b uint32) uint32 {
	if a > b {
		return a
	}
	return b
}

func writeRequest(w io.Writer, id uint64, method string, request proto.Message) error {
	// marshal request
	pbRequest := []byte{}
	if request != nil {
		var err error
		pbRequest, err = proto.Marshal(request)
		if err != nil {
			return err
		}
	}

	// compress serialized proto data
	compressedPbRequest := snappy.Encode(nil, pbRequest)

	// generate header
	header := &wire.RequestHeader{
		Id:                         id,
		Method:                     method,
		RawRequestLen:              uint32(len(pbRequest)),
		SnappyCompressedRequestLen: uint32(len(compressedPbRequest)),
		Checksum:                   crc32.ChecksumIEEE(compressedPbRequest),
	}

	// check header size
	pbHeader, err := proto.Marshal(header)
	if err != err {
		return err
	}
	if len(pbHeader) > int(wire.Const_MAX_REQUEST_HEADER_LEN) {
		return fmt.Errorf("protorpc.writeRequest: header larger than max_header_len: %d.", len(pbHeader))
	}

	// send header (more)
	if err := sendFrame(w, pbHeader); err != nil {
		return err
	}

	// send body (end)
	if err := sendFrame(w, compressedPbRequest); err != nil {
		return err
	}

	return nil
}

func readRequestHeader(r io.Reader, header *wire.RequestHeader) (err error) {
	// recv header (more)
	pbHeader, err := recvFrame(r, int(wire.Const_MAX_REQUEST_HEADER_LEN))
	if err != nil {
		return err
	}

	// Marshal Header
	err = proto.Unmarshal(pbHeader, header)
	if err != nil {
		return err
	}

	return nil
}

func readRequestBody(r io.Reader, header *wire.RequestHeader, request proto.Message) error {
	maxBodyLen := maxUint32(header.RawRequestLen, header.SnappyCompressedRequestLen)

	// recv body (end)
	compressedPbRequest, err := recvFrame(r, int(maxBodyLen))
	if err != nil {
		return err
	}

	// checksum
	if crc32.ChecksumIEEE(compressedPbRequest) != header.Checksum {
		return fmt.Errorf("protorpc.readRequestBody: unexpected checksum.")
	}

	// decode the compressed data
	pbRequest, err := snappy.Decode(nil, compressedPbRequest)
	if err != nil {
		return err
	}
	// check wire header: rawMsgLen
	if uint32(len(pbRequest)) != header.RawRequestLen {
		return fmt.Errorf("protorpc.readRequestBody: Unexcpeted header.RawRequestLen.")
	}

	// Unmarshal to proto message
	if request != nil {
		err = proto.Unmarshal(pbRequest, request)
		if err != nil {
			return err
		}
	}

	return nil
}

func writeResponse(w io.Writer, id uint64, serr string, response proto.Message) (err error) {
	// clean response if error
	if serr != "" {
		response = nil
	}

	// marshal response
	pbResponse := []byte{}
	if response != nil {
		pbResponse, err = proto.Marshal(response)
		if err != nil {
			return err
		}
	}

	// compress serialized proto data
	compressedPbResponse := snappy.Encode(nil, pbResponse)

	// generate header
	header := &wire.ResponseHeader{
		Id:                          id,
		Error:                       serr,
		RawResponseLen:              uint32(len(pbResponse)),
		SnappyCompressedResponseLen: uint32(len(compressedPbResponse)),
		Checksum:                    crc32.ChecksumIEEE(compressedPbResponse),
	}

	// check header size
	pbHeader, err := proto.Marshal(header)
	if err != err {
		return
	}

	// send header (more)
	if err = sendFrame(w, pbHeader); err != nil {
		return
	}

	// send body (end)
	if err = sendFrame(w, compressedPbResponse); err != nil {
		return
	}

	return nil
}

func readResponseHeader(r io.Reader, header *wire.ResponseHeader) error {
	// recv header (more)
	pbHeader, err := recvFrame(r, int(wire.Const_MAX_REQUEST_HEADER_LEN))
	if err != nil {
		return err
	}

	// Marshal Header
	err = proto.Unmarshal(pbHeader, header)
	if err != nil {
		return err
	}

	return nil
}

func readResponseBody(r io.Reader, header *wire.ResponseHeader, response proto.Message) error {
	maxBodyLen := int(maxUint32(header.RawResponseLen, header.SnappyCompressedResponseLen))

	// recv body (end)
	compressedPbResponse, err := recvFrame(r, maxBodyLen)
	if err != nil {
		return err
	}

	// checksum
	if crc32.ChecksumIEEE(compressedPbResponse) != header.Checksum {
		return fmt.Errorf("protorpc.readResponseBody: unexpected checksum.")
	}

	// decode the compressed data
	pbResponse, err := snappy.Decode(nil, compressedPbResponse)
	if err != nil {
		return err
	}
	// check wire header: rawMsgLen
	if uint32(len(pbResponse)) != header.RawResponseLen {
		return fmt.Errorf("protorpc.readResponseBody: Unexcpeted header.RawResponseLen.")
	}

	// Unmarshal to proto message
	if response != nil {
		err = proto.Unmarshal(pbResponse, response)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>add UseSappy and UseCrc32ChecksumIEEE<commit_after>// Copyright 2013 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package protorpc

import (
	"fmt"
	"hash/crc32"
	"io"

	wire "github.com/chai2010/protorpc/wire.pb"
	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
)

var (
	UseSappy             = true
	UseCrc32ChecksumIEEE = true
)

func maxUint32(a, b uint32) uint32 {
	if a > b {
		return a
	}
	return b
}

func writeRequest(w io.Writer, id uint64, method string, request proto.Message) error {
	// marshal request
	pbRequest := []byte{}
	if request != nil {
		var err error
		pbRequest, err = proto.Marshal(request)
		if err != nil {
			return err
		}
	}

	// compress serialized proto data
	compressedPbRequest := snappy.Encode(nil, pbRequest)

	// generate header
	header := &wire.RequestHeader{
		Id:                         id,
		Method:                     method,
		RawRequestLen:              uint32(len(pbRequest)),
		SnappyCompressedRequestLen: uint32(len(compressedPbRequest)),
		Checksum:                   crc32.ChecksumIEEE(compressedPbRequest),
	}

	if !UseSappy {
		header.SnappyCompressedRequestLen = 0
		compressedPbRequest = pbRequest
	}
	if !UseCrc32ChecksumIEEE {
		header.Checksum = 0
	}

	// check header size
	pbHeader, err := proto.Marshal(header)
	if err != err {
		return err
	}
	if len(pbHeader) > int(wire.Const_MAX_REQUEST_HEADER_LEN) {
		return fmt.Errorf("protorpc.writeRequest: header larger than max_header_len: %d.", len(pbHeader))
	}

	// send header (more)
	if err := sendFrame(w, pbHeader); err != nil {
		return err
	}

	// send body (end)
	if err := sendFrame(w, compressedPbRequest); err != nil {
		return err
	}

	return nil
}

func readRequestHeader(r io.Reader, header *wire.RequestHeader) (err error) {
	// recv header (more)
	pbHeader, err := recvFrame(r, int(wire.Const_MAX_REQUEST_HEADER_LEN))
	if err != nil {
		return err
	}

	// Marshal Header
	err = proto.Unmarshal(pbHeader, header)
	if err != nil {
		return err
	}

	return nil
}

func readRequestBody(r io.Reader, header *wire.RequestHeader, request proto.Message) error {
	maxBodyLen := maxUint32(header.RawRequestLen, header.SnappyCompressedRequestLen)

	// recv body (end)
	compressedPbRequest, err := recvFrame(r, int(maxBodyLen))
	if err != nil {
		return err
	}

	// checksum
	if header.Checksum != 0 {
		if crc32.ChecksumIEEE(compressedPbRequest) != header.Checksum {
			return fmt.Errorf("protorpc.readRequestBody: unexpected checksum.")
		}
	}

	var pbRequest []byte
	if header.SnappyCompressedRequestLen != 0 {
		// decode the compressed data
		pbRequest, err = snappy.Decode(nil, compressedPbRequest)
		if err != nil {
			return err
		}
		// check wire header: rawMsgLen
		if uint32(len(pbRequest)) != header.RawRequestLen {
			return fmt.Errorf("protorpc.readRequestBody: Unexcpeted header.RawRequestLen.")
		}
	} else {
		pbRequest = compressedPbRequest
	}

	// Unmarshal to proto message
	if request != nil {
		err = proto.Unmarshal(pbRequest, request)
		if err != nil {
			return err
		}
	}

	return nil
}

func writeResponse(w io.Writer, id uint64, serr string, response proto.Message) (err error) {
	// clean response if error
	if serr != "" {
		response = nil
	}

	// marshal response
	pbResponse := []byte{}
	if response != nil {
		pbResponse, err = proto.Marshal(response)
		if err != nil {
			return err
		}
	}

	// compress serialized proto data
	compressedPbResponse := snappy.Encode(nil, pbResponse)

	// generate header
	header := &wire.ResponseHeader{
		Id:                          id,
		Error:                       serr,
		RawResponseLen:              uint32(len(pbResponse)),
		SnappyCompressedResponseLen: uint32(len(compressedPbResponse)),
		Checksum:                    crc32.ChecksumIEEE(compressedPbResponse),
	}

	if !UseSappy {
		header.SnappyCompressedResponseLen = 0
		compressedPbResponse = pbResponse
	}
	if !UseCrc32ChecksumIEEE {
		header.Checksum = 0
	}

	// check header size
	pbHeader, err := proto.Marshal(header)
	if err != err {
		return
	}

	// send header (more)
	if err = sendFrame(w, pbHeader); err != nil {
		return
	}

	// send body (end)
	if err = sendFrame(w, compressedPbResponse); err != nil {
		return
	}

	return nil
}

func readResponseHeader(r io.Reader, header *wire.ResponseHeader) error {
	// recv header (more)
	pbHeader, err := recvFrame(r, int(wire.Const_MAX_REQUEST_HEADER_LEN))
	if err != nil {
		return err
	}

	// Marshal Header
	err = proto.Unmarshal(pbHeader, header)
	if err != nil {
		return err
	}

	return nil
}

func readResponseBody(r io.Reader, header *wire.ResponseHeader, response proto.Message) error {
	maxBodyLen := int(maxUint32(header.RawResponseLen, header.SnappyCompressedResponseLen))

	// recv body (end)
	compressedPbResponse, err := recvFrame(r, maxBodyLen)
	if err != nil {
		return err
	}

	// checksum
	if header.Checksum != 0 {
		if crc32.ChecksumIEEE(compressedPbResponse) != header.Checksum {
			return fmt.Errorf("protorpc.readResponseBody: unexpected checksum.")
		}
	}

	var pbResponse []byte
	if header.SnappyCompressedResponseLen != 0 {
		// decode the compressed data
		pbResponse, err = snappy.Decode(nil, compressedPbResponse)
		if err != nil {
			return err
		}
		// check wire header: rawMsgLen
		if uint32(len(pbResponse)) != header.RawResponseLen {
			return fmt.Errorf("protorpc.readResponseBody: Unexcpeted header.RawResponseLen.")
		}
	} else {
		pbResponse = compressedPbResponse
	}

	// Unmarshal to proto message
	if response != nil {
		err = proto.Unmarshal(pbResponse, response)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Ardan Studios. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE handle.

// Package work manages a pool of routines to perform work.
package work

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 2
)

// ErrorInvalidMinRoutines is the error for the invalid minRoutine parameter.
var ErrorInvalidMinRoutines = errors.New("Invalid minimum number of routines")

// ErrorInvalidIdleTime is the error for the invalid idle time parameter.
var ErrorInvalidIdleTime = errors.New("Invalid duration for idle time")

// ErrorInvalidStatTime is the error for the invalid stat time parameter.
var ErrorInvalidStatTime = errors.New("Invalid duration for stat time")

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(id int)
}

// Work provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Work struct {
	minRoutines int            // Minumum number of routines always in the pool.
	idleTime    time.Duration  // Time for routines to die on idle.
	statTime    time.Duration  // Time to display stats.
	counter     int            // Maintains a running total number of routines ever created.
	tasks       chan Worker    // Unbuffered channel that work is sent into.
	control     chan int       // Unbuffered channel that work for the manager is send into.
	kill        chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown    chan struct{}  // Closed when the Work pool is being shutdown.
	wg          sync.WaitGroup // Manages the number of routines for shutdown.
	routines    int64          // Number of routines
	active      int64          // Active number of routines in the work pool.
	pending     int64          // Pending number of routines waiting to submit work.
}

// manager controls changes to the work pool including stats
// and shutting down.
func (w *Work) manager() {
	w.wg.Add(1)

	go func() {
		log.Println("Work : manager : Started")

		// Create a timer to run stats.
		var stats <-chan time.Time
		timer := time.NewTimer(w.statTime)
		stats = timer.C

		for {
			select {
			case <-w.shutdown:
				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&w.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					w.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				w.wg.Done()
				return

			case c := <-w.control:
				switch c {
				case addRoutine:
					log.Println("Work : manager : Info : Add Routine")

					// Capture a unique id.
					w.counter++

					// Add to the counts.
					w.wg.Add(1)
					atomic.AddInt64(&w.routines, 1)

					// Create the routine.
					go w.work(w.counter)

				case rmvRoutine:
					log.Println("Work : manager : Info : Remove Routine")

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&w.routines))

					// Are there routines to remove.
					if routines <= w.minRoutines {
						log.Println("Work : manager : Info : Remove Routine Cancelled")
						break
					}

					// Send a kill signal to remove a routine.
					w.kill <- struct{}{}
				}

			case <-stats:
				// Capture the stats.
				routines := atomic.LoadInt64(&w.routines)
				pending := atomic.LoadInt64(&w.pending)
				active := atomic.LoadInt64(&w.active)

				// Display the stats.
				fmt.Printf("Work : manager : Stats : G[%d] P[%d] A[%d]\n", routines, pending, active)

				// Reset the clock.
				timer.Reset(w.statTime)
			}
		}
	}()
}

// New creates a new Worker.
func New(minRoutines int, idleTime time.Duration, statTime time.Duration) (*Work, error) {
	if minRoutines < 0 {
		return nil, ErrorInvalidMinRoutines
	}

	if idleTime < time.Minute {
		return nil, ErrorInvalidIdleTime
	}

	if statTime < time.Millisecond {
		return nil, ErrorInvalidStatTime
	}

	w := Work{
		minRoutines: minRoutines,
		idleTime:    idleTime,
		statTime:    statTime,
		tasks:       make(chan Worker),
		control:     make(chan int),
		kill:        make(chan struct{}),
		shutdown:    make(chan struct{}),
	}

	// Start the manager.
	w.manager()

	// Add the routines.
	w.Add(minRoutines)

	return &w, nil
}

// Add creates routines to process work or sets a count for
// routines to terminate.
func (w *Work) Add(routines int) {
	if routines == 0 {
		return
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	for i := 0; i < routines; i++ {
		w.control <- cmd
	}
}

// work performs the users work and keeps stats.
func (w *Work) work(id int) {
	// Create a timer to track idle time.
	var idle <-chan time.Time
	var timer *time.Timer

	// Set the timer for routines about the min mark.
	if id > w.minRoutines {
		timer = time.NewTimer(w.idleTime)
		idle = timer.C
	}

done:
	for {
		select {
		case t := <-w.tasks:
			atomic.AddInt64(&w.active, 1)
			{
				// Perform the work.
				t.Work(id)
			}
			atomic.AddInt64(&w.active, -1)

		case <-w.kill:
			break done

		case <-idle:
			break done
		}

		// If this goroutine can die on idle time
		// then reset the timer.
		if timer != nil {
			timer.Reset(w.idleTime)
		}
	}

	// Decrement the counts.
	atomic.AddInt64(&w.routines, -1)
	w.wg.Done()

	log.Println("Work : gr : Info : Shutdown")
}

// Run wait for the goroutine pool to take the work
// to be executed.
func (w *Work) Run(work Worker) {
	atomic.AddInt64(&w.pending, 1)
	{
		w.tasks <- work
	}
	atomic.AddInt64(&w.pending, -1)
}

// Shutdown waits for all the workers to finish.
func (w *Work) Shutdown() {
	close(w.shutdown)
	w.wg.Wait()
}
<commit_msg>Removed idle time. Goroutines are cheap and this is un-necessary.<commit_after>// Copyright 2014 Ardan Studios. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE handle.

// Package work manages a pool of routines to perform work.
package work

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 2
)

// ErrorInvalidMinRoutines is the error for the invalid minRoutine parameter.
var ErrorInvalidMinRoutines = errors.New("Invalid minimum number of routines")

// ErrorInvalidStatTime is the error for the invalid stat time parameter.
var ErrorInvalidStatTime = errors.New("Invalid duration for stat time")

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(id int)
}

// Work provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Work struct {
	minRoutines int            // Minumum number of routines always in the pool.
	statTime    time.Duration  // Time to display stats.
	counter     int            // Maintains a running total number of routines ever created.
	tasks       chan Worker    // Unbuffered channel that work is sent into.
	control     chan int       // Unbuffered channel that work for the manager is send into.
	kill        chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown    chan struct{}  // Closed when the Work pool is being shutdown.
	wg          sync.WaitGroup // Manages the number of routines for shutdown.
	routines    int64          // Number of routines
	active      int64          // Active number of routines in the work pool.
	pending     int64          // Pending number of routines waiting to submit work.
}

// New creates a new Worker.
func New(minRoutines int, statTime time.Duration) (*Work, error) {
	if minRoutines < 0 {
		return nil, ErrorInvalidMinRoutines
	}

	if statTime < time.Millisecond {
		return nil, ErrorInvalidStatTime
	}

	w := Work{
		minRoutines: minRoutines,
		statTime:    statTime,
		tasks:       make(chan Worker),
		control:     make(chan int),
		kill:        make(chan struct{}),
		shutdown:    make(chan struct{}),
	}

	// Start the manager.
	w.manager()

	// Add the routines.
	w.Add(minRoutines)

	return &w, nil
}

// Add creates routines to process work or sets a count for
// routines to terminate.
func (w *Work) Add(routines int) {
	if routines == 0 {
		return
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	for i := 0; i < routines; i++ {
		w.control <- cmd
	}
}

// work performs the users work and keeps stats.
func (w *Work) work(id int) {
done:
	for {
		select {
		case t := <-w.tasks:
			atomic.AddInt64(&w.active, 1)
			{
				// Perform the work.
				t.Work(id)
			}
			atomic.AddInt64(&w.active, -1)

		case <-w.kill:
			break done
		}
	}

	// Decrement the counts.
	atomic.AddInt64(&w.routines, -1)
	w.wg.Done()

	log.Println("Work : gr : Info : Shutdown")
}

// Run wait for the goroutine pool to take the work
// to be executed.
func (w *Work) Run(work Worker) {
	atomic.AddInt64(&w.pending, 1)
	{
		w.tasks <- work
	}
	atomic.AddInt64(&w.pending, -1)
}

// Shutdown waits for all the workers to finish.
func (w *Work) Shutdown() {
	close(w.shutdown)
	w.wg.Wait()
}

// manager controls changes to the work pool including stats
// and shutting down.
func (w *Work) manager() {
	w.wg.Add(1)

	go func() {
		log.Println("Work : manager : Started")

		// Create a timer to run stats.
		timer := time.NewTimer(w.statTime)

		for {
			select {
			case <-w.shutdown:
				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&w.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					w.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				w.wg.Done()
				return

			case c := <-w.control:
				switch c {
				case addRoutine:
					log.Println("Work : manager : Info : Add Routine")

					// Capture a unique id.
					w.counter++

					// Add to the counts.
					w.wg.Add(1)
					atomic.AddInt64(&w.routines, 1)

					// Create the routine.
					go w.work(w.counter)

				case rmvRoutine:
					log.Println("Work : manager : Info : Remove Routine")

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&w.routines))

					// Are there routines to remove.
					if routines <= w.minRoutines {
						log.Println("Work : manager : Info : Remove Routine Cancelled")
						break
					}

					// Send a kill signal to remove a routine.
					w.kill <- struct{}{}
				}

			case <-timer.C:
				// Capture the stats.
				routines := atomic.LoadInt64(&w.routines)
				pending := atomic.LoadInt64(&w.pending)
				active := atomic.LoadInt64(&w.active)

				// Display the stats.
				fmt.Printf("Work : manager : Stats : G[%d] P[%d] A[%d]\n", routines, pending, active)

				// Reset the clock.
				timer.Reset(w.statTime)
			}
		}
	}()
}
<|endoftext|>


<commit_before>/*
The MIT License (MIT)

Copyright (c) 2015 John Ko

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"bytes"
	"fmt"
	"github.com/gorilla/mux"
	"io"
	//"mime"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

func headHandler(w http.ResponseWriter, r *http.Request) {
	// force close because curl tries to keep open
	w.Header().Set("Connection", "close")
	if allowedHead() {
		vars := mux.Vars(r)
		hash := vars["hash"]
		_, _, _, err := storage.Head(hash)
		//filename, contentLength, _, err := storage.Head(hash)
		//contentType := mime.TypeByExtension(filepath.Ext(filename))
		if err != nil {
			// we only care about the status code
			w.WriteHeader(404)
			return
		} else {
			// you may want the json?
			//fmt.Fprintf(w, "{\"sha512\":\"%s\",\"filename\":\"%s\",\"length\":%d,\"content_type\":\"%s\",\"stub\":true}", hash, filename, contentLength, strings.Split(contentType, ";")[0])
		}
	} else {
		// we only care about the status code
		w.WriteHeader(403)
	}
}

func SplitHashToPairSlash(token string) string {
	// split hash to pairs because a folder of 400+ items is slow
	pairs := regexp.MustCompile("[0-9a-f]{2}").FindAll([]byte(token), -1)
	// join pairs with os.PathSeparator
	newtoken := bytes.Join(pairs, []byte(string(os.PathSeparator)))
	return string(newtoken[:])
}

func NameLengthTime(path string) (filename string, contentLength uint64, modTime time.Time, err error) {
	// content length
	var fi os.FileInfo
	if fi, err = os.Lstat(filepath.Join(path, "data")); err != nil {
		return
	}
	contentLength = uint64(fi.Size())
	modTime = fi.ModTime()
	// Use tail to get the last real filename
	var lastname []byte
	lastname, _ = exec.Command(cmdTAIL, "-n", "1", filepath.Join(path, "filename")).Output()
	// Assume the first output before space is the mimeType
	filename = strings.TrimSpace(fmt.Sprintf("%s", lastname))
	return
}

func Sha512(str string, word string) (hash string, err error) {
	// FreeBSD specific call /sbin/sha512 instead of using the import crypto/sha512
	// because the import has high memory usage (loads the data in RAM)
	// and Go lang uses garbage collection so the high RAM lingers
	// Assume the output is the hash, need to trim \n

	//hash, err := exec.Command(cmdSHA512, "-q", "-s", word).Output()
	//if err != nil {
	//      return
	//}

	// TODO: is shasum more universal on *nix systems?
	cmd := exec.Command(cmdSHASUM, "--algorithm", "512", str)
	if word != "" {
		cmd.Stdin = strings.NewReader(word)
	}
	tmpout, err := cmd.Output()
	// Assume the first output before space is the hash
	hash = strings.Split(strings.TrimSpace(fmt.Sprintf("%s", tmpout)), " ")[0]
	return
}

func Sha512Word(word string) (hash string, err error) {
	hash, err = Sha512("-", word)
	if err != nil {
		return
	}
	return
}

func (s *LocalStorage) saveFilename(hash string, filename string) {
	var err error
	newpath := filepath.Join(s.basedir, SplitHashToPairSlash(hash))
	var f1 io.WriteCloser
	f1, err = os.OpenFile(filepath.Join(newpath, "filename"), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
	if err == nil {
		defer f1.Close()
		io.Copy(f1, strings.NewReader(fmt.Sprintf("%s\n", filename)))
	}
}

func (s *LocalStorage) HardLinkSha512Path(oldpath string, filename string) (hash string, contentLength uint64, err error) {
	var fi os.FileInfo
	if fi, err = os.Lstat(oldpath); err != nil {
		return
	}
	contentLength = uint64(fi.Size())
	hash, err = Sha512(oldpath, "")
	if err != nil {
		return
	}
	newpath := filepath.Join(s.basedir, SplitHashToPairSlash(hash))
	// mkdir -p
	if err = os.MkdirAll(newpath, 0700); err != nil && !os.IsExist(err) {
		return
	}
	// Link the oldpath to sha512/data
	if err = os.Link(oldpath, filepath.Join(newpath, "data")); err != nil {
		if strings.Index(err.Error(), "file exists") >= 0 {
			err = nil
		}
	}
	storage.saveFilename(newpath, filename)
	os.Remove(oldpath)
	return
}

func (s *LocalStorage) HardLinkSha512(token string, filename string) (hash string, contentLength uint64, err error) {
	oldpath := filepath.Join(config.Temp, token, filename)
	hash, contentLength, err = storage.HardLinkSha512Path(oldpath, filename)
	return
}

func (s *LocalStorage) DeleteFile(token string, filename string) error {
	oldpath := filepath.Join(config.Temp, token)
	os.Remove(filepath.Join(oldpath, filename))
	os.Remove(oldpath)
	return nil
}
<commit_msg>storage.saveFilename(hash, filename)<commit_after>/*
The MIT License (MIT)

Copyright (c) 2015 John Ko

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"bytes"
	"fmt"
	"github.com/gorilla/mux"
	"io"
	//"mime"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

func headHandler(w http.ResponseWriter, r *http.Request) {
	// force close because curl tries to keep open
	w.Header().Set("Connection", "close")
	if allowedHead() {
		vars := mux.Vars(r)
		hash := vars["hash"]
		_, _, _, err := storage.Head(hash)
		//filename, contentLength, _, err := storage.Head(hash)
		//contentType := mime.TypeByExtension(filepath.Ext(filename))
		if err != nil {
			// we only care about the status code
			w.WriteHeader(404)
			return
		} else {
			// you may want the json?
			//fmt.Fprintf(w, "{\"sha512\":\"%s\",\"filename\":\"%s\",\"length\":%d,\"content_type\":\"%s\",\"stub\":true}", hash, filename, contentLength, strings.Split(contentType, ";")[0])
		}
	} else {
		// we only care about the status code
		w.WriteHeader(403)
	}
}

func SplitHashToPairSlash(token string) string {
	// split hash to pairs because a folder of 400+ items is slow
	pairs := regexp.MustCompile("[0-9a-f]{2}").FindAll([]byte(token), -1)
	// join pairs with os.PathSeparator
	newtoken := bytes.Join(pairs, []byte(string(os.PathSeparator)))
	return string(newtoken[:])
}

func NameLengthTime(path string) (filename string, contentLength uint64, modTime time.Time, err error) {
	// content length
	var fi os.FileInfo
	if fi, err = os.Lstat(filepath.Join(path, "data")); err != nil {
		return
	}
	contentLength = uint64(fi.Size())
	modTime = fi.ModTime()
	// Use tail to get the last real filename
	var lastname []byte
	lastname, _ = exec.Command(cmdTAIL, "-n", "1", filepath.Join(path, "filename")).Output()
	// Assume the first output before space is the mimeType
	filename = strings.TrimSpace(fmt.Sprintf("%s", lastname))
	return
}

func Sha512(str string, word string) (hash string, err error) {
	// FreeBSD specific call /sbin/sha512 instead of using the import crypto/sha512
	// because the import has high memory usage (loads the data in RAM)
	// and Go lang uses garbage collection so the high RAM lingers
	// Assume the output is the hash, need to trim \n

	//hash, err := exec.Command(cmdSHA512, "-q", "-s", word).Output()
	//if err != nil {
	//      return
	//}

	// TODO: is shasum more universal on *nix systems?
	cmd := exec.Command(cmdSHASUM, "--algorithm", "512", str)
	if word != "" {
		cmd.Stdin = strings.NewReader(word)
	}
	tmpout, err := cmd.Output()
	// Assume the first output before space is the hash
	hash = strings.Split(strings.TrimSpace(fmt.Sprintf("%s", tmpout)), " ")[0]
	return
}

func Sha512Word(word string) (hash string, err error) {
	hash, err = Sha512("-", word)
	if err != nil {
		return
	}
	return
}

func (s *LocalStorage) saveFilename(hash string, filename string) {
	var err error
	newpath := filepath.Join(s.basedir, SplitHashToPairSlash(hash))
	var f1 io.WriteCloser
	f1, err = os.OpenFile(filepath.Join(newpath, "filename"), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
	if err == nil {
		defer f1.Close()
		io.Copy(f1, strings.NewReader(fmt.Sprintf("%s\n", filename)))
	}
}

func (s *LocalStorage) HardLinkSha512Path(oldpath string, filename string) (hash string, contentLength uint64, err error) {
	var fi os.FileInfo
	if fi, err = os.Lstat(oldpath); err != nil {
		return
	}
	contentLength = uint64(fi.Size())
	hash, err = Sha512(oldpath, "")
	if err != nil {
		return
	}
	newpath := filepath.Join(s.basedir, SplitHashToPairSlash(hash))
	// mkdir -p
	if err = os.MkdirAll(newpath, 0700); err != nil && !os.IsExist(err) {
		return
	}
	// Link the oldpath to sha512/data
	if err = os.Link(oldpath, filepath.Join(newpath, "data")); err != nil {
		if strings.Index(err.Error(), "file exists") >= 0 {
			err = nil
		}
	}
	storage.saveFilename(hash, filename)
	os.Remove(oldpath)
	return
}

func (s *LocalStorage) HardLinkSha512(token string, filename string) (hash string, contentLength uint64, err error) {
	oldpath := filepath.Join(config.Temp, token, filename)
	hash, contentLength, err = storage.HardLinkSha512Path(oldpath, filename)
	return
}

func (s *LocalStorage) DeleteFile(token string, filename string) error {
	oldpath := filepath.Join(config.Temp, token)
	os.Remove(filepath.Join(oldpath, filename))
	os.Remove(oldpath)
	return nil
}
<|endoftext|>


<commit_before>package dbsplitter

import (
	"fmt"
	"github.com/go-xorm/xorm"
)

type DbSplitterXorm struct{}

func (d *DbSplitterXorm) GetEngines(dbs []string) (map[string]*xorm.Engine, error) {
	engines := map[string]*xorm.Engine{}
	for _, db := range dbs {
		params := getConnectionStrings(db)
		engines[db] = d.getEngine(params)
	}
	return engines, nil
}

func (d *DbSplitterXorm) getEngine(params map[string]string) *xorm.Engine {
	dsn := d.parseParameter(params)
	db, err := xorm.NewEngine("mysql", dsn)
	if err != nil {
		panic("cannot get engines. params=" + dsn)
	}
	return db
}

func (d *DbSplitterXorm) parseParameter(params map[string]string) string {
	return fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8", params["user"], params["pass"], params["host"], params["port"], params["dbname"])
}
<commit_msg>Change xorm naming mapper to Gonic<commit_after>package dbsplitter

import (
	"fmt"

	"github.com/go-xorm/xorm"
	"github.com/go-xorm/core"
)

type DbSplitterXorm struct{}

func (d *DbSplitterXorm) GetEngines(dbs []string) (map[string]*xorm.Engine, error) {
	engines := map[string]*xorm.Engine{}
	for _, db := range dbs {
		params := getConnectionStrings(db)
		engines[db] = d.getEngine(params)
	}
	return engines, nil
}

func (d *DbSplitterXorm) getEngine(params map[string]string) *xorm.Engine {
	dsn := d.parseParameter(params)
	db, err := xorm.NewEngine("mysql", dsn)
	if err != nil {
		panic("cannot get engines. params=" + dsn)
	}
	db.SetMapper(core.NewCacheMapper(new(core.GonicMapper)))
	return db
}

func (d *DbSplitterXorm) parseParameter(params map[string]string) string {
	return fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8", params["user"], params["pass"], params["host"], params["port"], params["dbname"])
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/crooks/yamn/config"
	"github.com/crooks/yamn/idlog"
	"github.com/crooks/yamn/keymgr"
	"github.com/luksen/maildir"
)

const (
	version        string = "0.2.5"
	dayLength      int    = 24 * 60 * 60 // Day in seconds
	maxFragLength         = 17910
	maxCopies             = 5
	base64LineWrap        = 64
	rfc5322date           = "Mon, 2 Jan 2006 15:04:05 -0700"
	shortdate             = "2 Jan 2006"
)

var (
	// flags - Command line flags
	flags *config.Flags
	// cfg - Config parameters
	cfg *config.Config
	// Trace loglevel
	Trace *log.Logger
	// Info loglevel
	Info *log.Logger
	// Warn loglevel
	Warn *log.Logger
	// Error loglevel
	Error *log.Logger
	// Pubring - Public Keyring
	Pubring *keymgr.Pubring
	// IDDb - Message ID log (replay protection)
	IDDb *idlog.IDLog
	// ChunkDb - Chunk database
	ChunkDb *Chunk
)

func logInit(
	traceHandle io.Writer,
	infoHandle io.Writer,
	warnHandle io.Writer,
	errorHandle io.Writer) {

	Trace = log.New(traceHandle,
		"Trace: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Info = log.New(infoHandle,
		"Info: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Warn = log.New(warnHandle,
		"Warn: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Error = log.New(errorHandle,
		"Error: ",
		log.Ldate|log.Ltime|log.Lshortfile)
}

func main() {
	var err error
	flags = config.ParseFlags()
	// Some config defaults are derived from flags so ParseConfig is a flags method
	cfg, err = flags.ParseConfig()
	if err != nil {
		// No logging is defined at this point so log the error to stderr
		fmt.Fprintf(os.Stderr, "Unable to parse config file: %v", err)
		os.Exit(1)
	}
	if cfg.General.LogToFile {
		logfile, err := os.OpenFile(
			cfg.Files.Logfile,
			os.O_RDWR|os.O_CREATE|os.O_APPEND,
			0640,
		)
		if err != nil {
			fmt.Fprintf(
				os.Stderr,
				"Error opening logfile: %s.\n",
				err,
			)
			os.Exit(1)
		}
		switch strings.ToLower(cfg.General.Loglevel) {
		case "trace":
			logInit(logfile,
				logfile,
				logfile,
				logfile,
			)
		case "info":
			logInit(ioutil.Discard,
				logfile,
				logfile,
				logfile,
			)
		case "warn":
			logInit(ioutil.Discard,
				ioutil.Discard,
				logfile,
				logfile,
			)
		case "error":
			logInit(ioutil.Discard,
				ioutil.Discard,
				ioutil.Discard,
				logfile,
			)
		default:
			fmt.Fprintf(
				os.Stderr,
				"Unknown loglevel: %s.  Assuming \"Info\".\n",
				cfg.General.Loglevel,
			)
			logInit(ioutil.Discard,
				logfile,
				logfile,
				logfile,
			)
		}
	} else {
		switch strings.ToLower(cfg.General.Loglevel) {
		case "trace":
			logInit(os.Stdout,
				os.Stdout,
				os.Stdout,
				os.Stderr,
			)
		case "info":
			logInit(ioutil.Discard,
				os.Stdout,
				os.Stdout,
				os.Stderr,
			)
		case "warn":
			logInit(ioutil.Discard,
				ioutil.Discard,
				os.Stdout,
				os.Stderr,
			)
		case "error":
			logInit(
				ioutil.Discard,
				ioutil.Discard,
				ioutil.Discard,
				os.Stderr,
			)
		default:
			fmt.Fprintf(
				os.Stderr,
				"Unknown loglevel: %s.  Assuming \"Info\".\n",
				cfg.General.Loglevel,
			)
			logInit(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)
		} // End of stdout/stderr logging setup
	} // End of logging setup

	// Inform the user which (if any) config file was used.
	if cfg.Files.Config != "" {
		Info.Printf("Using config file: %s", cfg.Files.Config)
	} else {
		Warn.Println("No config file was found. Resorting to defaults")
	}

	// If the debug flag is set, print the config in JSON format and then exit.
	if flags.Debug {
		y, err := cfg.Debug()
		if err != nil {
			fmt.Printf("Debugging Error: %s\n", err)
			os.Exit(1)
		}
		fmt.Printf("%s\n", y)
		os.Exit(0)
	}

	if flags.Client {
		mixprep()
	} else if flags.Stdin {
		dir := maildir.Dir(cfg.Files.Maildir)
		newmsg, err := dir.NewDelivery()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		newmsg.Write(stdin)
		newmsg.Close()
	} else if flags.Remailer {
		err = loopServer()
		if err != nil {
			panic(err)
		}
	} else if flags.Dummy {
		injectDummy()
	} else if flags.Refresh {
		fmt.Printf("Keyring refresh: from=%s, to=%s\n", cfg.Urls.Pubring, cfg.Files.Pubring)
		httpGet(cfg.Urls.Pubring, cfg.Files.Pubring)
		fmt.Printf("Stats refresh: from=%s, to=%s\n", cfg.Urls.Mlist2, cfg.Files.Mlist2)
		httpGet(cfg.Urls.Mlist2, cfg.Files.Mlist2)
	}
	if flags.Send {
		// Flush the outbound pool
		poolOutboundSend()
	}
}
<commit_msg>Do debugging in YAML and fix the version flag<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/crooks/yamn/config"
	"github.com/crooks/yamn/idlog"
	"github.com/crooks/yamn/keymgr"
	"github.com/luksen/maildir"
)

const (
	version        string = "0.2.5"
	dayLength      int    = 24 * 60 * 60 // Day in seconds
	maxFragLength         = 17910
	maxCopies             = 5
	base64LineWrap        = 64
	rfc5322date           = "Mon, 2 Jan 2006 15:04:05 -0700"
	shortdate             = "2 Jan 2006"
)

var (
	// flags - Command line flags
	flags *config.Flags
	// cfg - Config parameters
	cfg *config.Config
	// Trace loglevel
	Trace *log.Logger
	// Info loglevel
	Info *log.Logger
	// Warn loglevel
	Warn *log.Logger
	// Error loglevel
	Error *log.Logger
	// Pubring - Public Keyring
	Pubring *keymgr.Pubring
	// IDDb - Message ID log (replay protection)
	IDDb *idlog.IDLog
	// ChunkDb - Chunk database
	ChunkDb *Chunk
)

func logInit(
	traceHandle io.Writer,
	infoHandle io.Writer,
	warnHandle io.Writer,
	errorHandle io.Writer) {

	Trace = log.New(traceHandle,
		"Trace: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Info = log.New(infoHandle,
		"Info: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Warn = log.New(warnHandle,
		"Warn: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Error = log.New(errorHandle,
		"Error: ",
		log.Ldate|log.Ltime|log.Lshortfile)
}

func main() {
	var err error
	flags = config.ParseFlags()
	if flags.Version {
		fmt.Println(version)
		os.Exit(0)
	}
	// Some config defaults are derived from flags so ParseConfig is a flags method
	cfg, err = flags.ParseConfig()
	if err != nil {
		// No logging is defined at this point so log the error to stderr
		fmt.Fprintf(os.Stderr, "Unable to parse config file: %v", err)
		os.Exit(1)
	}
	// If the debug flag is set, print the config and exit
	if flags.Debug {
		y, err := cfg.Debug()
		if err != nil {
			fmt.Printf("Debugging Error: %s\n", err)
			os.Exit(1)
		}
		fmt.Printf("%s\n", y)
		os.Exit(0)
	}

	// Set up some logging based on loglevel and LogToFile
	if cfg.General.LogToFile {
		logfile, err := os.OpenFile(
			cfg.Files.Logfile,
			os.O_RDWR|os.O_CREATE|os.O_APPEND,
			0640,
		)
		if err != nil {
			fmt.Fprintf(
				os.Stderr,
				"Error opening logfile: %s.\n",
				err,
			)
			os.Exit(1)
		}
		switch strings.ToLower(cfg.General.Loglevel) {
		case "trace":
			logInit(logfile,
				logfile,
				logfile,
				logfile,
			)
		case "info":
			logInit(ioutil.Discard,
				logfile,
				logfile,
				logfile,
			)
		case "warn":
			logInit(ioutil.Discard,
				ioutil.Discard,
				logfile,
				logfile,
			)
		case "error":
			logInit(ioutil.Discard,
				ioutil.Discard,
				ioutil.Discard,
				logfile,
			)
		default:
			fmt.Fprintf(
				os.Stderr,
				"Unknown loglevel: %s.  Assuming \"Info\".\n",
				cfg.General.Loglevel,
			)
			logInit(ioutil.Discard,
				logfile,
				logfile,
				logfile,
			)
		}
	} else {
		switch strings.ToLower(cfg.General.Loglevel) {
		case "trace":
			logInit(os.Stdout,
				os.Stdout,
				os.Stdout,
				os.Stderr,
			)
		case "info":
			logInit(ioutil.Discard,
				os.Stdout,
				os.Stdout,
				os.Stderr,
			)
		case "warn":
			logInit(ioutil.Discard,
				ioutil.Discard,
				os.Stdout,
				os.Stderr,
			)
		case "error":
			logInit(
				ioutil.Discard,
				ioutil.Discard,
				ioutil.Discard,
				os.Stderr,
			)
		default:
			fmt.Fprintf(
				os.Stderr,
				"Unknown loglevel: %s.  Assuming \"Info\".\n",
				cfg.General.Loglevel,
			)
			logInit(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)
		} // End of stdout/stderr logging setup
	} // End of logging setup

	// Inform the user which (if any) config file was used.
	if cfg.Files.Config != "" {
		Info.Printf("Using config file: %s", cfg.Files.Config)
	} else {
		Warn.Println("No config file was found. Resorting to defaults")
	}

	if flags.Client {
		mixprep()
	} else if flags.Stdin {
		dir := maildir.Dir(cfg.Files.Maildir)
		newmsg, err := dir.NewDelivery()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		newmsg.Write(stdin)
		newmsg.Close()
	} else if flags.Remailer {
		err = loopServer()
		if err != nil {
			panic(err)
		}
	} else if flags.Dummy {
		injectDummy()
	} else if flags.Refresh {
		fmt.Printf("Keyring refresh: from=%s, to=%s\n", cfg.Urls.Pubring, cfg.Files.Pubring)
		httpGet(cfg.Urls.Pubring, cfg.Files.Pubring)
		fmt.Printf("Stats refresh: from=%s, to=%s\n", cfg.Urls.Mlist2, cfg.Files.Mlist2)
		httpGet(cfg.Urls.Mlist2, cfg.Files.Mlist2)
	}
	if flags.Send {
		// Flush the outbound pool
		poolOutboundSend()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	stdlog "log"
	"os"

	"github.com/Masterminds/log-go"
	"github.com/crooks/jlog"
	"github.com/crooks/yamn/config"
	"github.com/crooks/yamn/idlog"
	"github.com/crooks/yamn/keymgr"
	"github.com/luksen/maildir"
)

const (
	version        string = "0.2.5"
	dayLength      int    = 24 * 60 * 60 // Day in seconds
	maxFragLength         = 17910
	maxCopies             = 5
	base64LineWrap        = 64
	rfc5322date           = "Mon, 2 Jan 2006 15:04:05 -0700"
	shortdate             = "2 Jan 2006"
)

var (
	// flags - Command line flags
	flag *config.Flags
	// cfg - Config parameters
	cfg *config.Config
	// Pubring - Public Keyring
	Pubring *keymgr.Pubring
	// IDDb - Message ID log (replay protection)
	IDDb *idlog.IDLog
	// ChunkDb - Chunk database
	ChunkDb *Chunk
)

func main() {
	var err error
	flag, cfg = config.GetCfg()
	if flag.Version {
		fmt.Println(version)
		os.Exit(0)
	}
	// If the debug flag is set, print the config and exit
	if flag.Debug {
		y, err := cfg.Debug()
		if err != nil {
			fmt.Printf("Debugging Error: %s\n", err)
			os.Exit(1)
		}
		fmt.Printf("%s\n", y)
		os.Exit(0)
	}

	// Set up logging
	loglevel, err := log.Atoi(cfg.General.Loglevel)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: Unknown loglevel", cfg.General.Loglevel)
		os.Exit(1)
	}
	// If we're logging to a file, open the file and redirect output to it
	if cfg.General.LogToFile && cfg.General.LogToJournal {
		fmt.Fprintln(os.Stderr, "Cannot log to file and journal")
		os.Exit(1)
	} else if cfg.General.LogToFile {
		logfile, err := os.OpenFile(cfg.Files.Logfile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%s: Error opening logfile: %v", cfg.Files.Logfile, err)
			os.Exit(1)
		}
		stdlog.SetOutput(logfile)
		log.Current = log.StdLogger{Level: loglevel}
	} else if cfg.General.LogToJournal {
		log.Current = jlog.NewJournal(loglevel)
	} else {
		log.Current = log.StdLogger{Level: loglevel}
	}

	// Inform the user which (if any) config file was used.
	if cfg.Files.Config != "" {
		log.Infof("Using config file: %s", cfg.Files.Config)
	} else {
		log.Warn("No config file was found. Resorting to defaults")
	}

	// Setup complete, time to do some work
	if flag.Client {
		mixprep()
	} else if flag.Stdin {
		dir := maildir.Dir(cfg.Files.Maildir)
		newmsg, err := dir.NewDelivery()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		newmsg.Write(stdin)
		newmsg.Close()
	} else if flag.Remailer {
		err = loopServer()
		if err != nil {
			panic(err)
		}
	} else if flag.Dummy {
		injectDummy()
	} else if flag.Refresh {
		fmt.Printf("Keyring refresh: from=%s, to=%s\n", cfg.Urls.Pubring, cfg.Files.Pubring)
		httpGet(cfg.Urls.Pubring, cfg.Files.Pubring)
		fmt.Printf("Stats refresh: from=%s, to=%s\n", cfg.Urls.Mlist2, cfg.Files.Mlist2)
		httpGet(cfg.Urls.Mlist2, cfg.Files.Mlist2)
	}
	if flag.Send {
		// Flush the outbound pool
		poolOutboundSend()
	}
}
<commit_msg>Bump version number<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	stdlog "log"
	"os"

	"github.com/Masterminds/log-go"
	"github.com/crooks/jlog"
	"github.com/crooks/yamn/config"
	"github.com/crooks/yamn/idlog"
	"github.com/crooks/yamn/keymgr"
	"github.com/luksen/maildir"
)

const (
	version        string = "0.2.6"
	dayLength      int    = 24 * 60 * 60 // Day in seconds
	maxFragLength         = 17910
	maxCopies             = 5
	base64LineWrap        = 64
	rfc5322date           = "Mon, 2 Jan 2006 15:04:05 -0700"
	shortdate             = "2 Jan 2006"
)

var (
	// flags - Command line flags
	flag *config.Flags
	// cfg - Config parameters
	cfg *config.Config
	// Pubring - Public Keyring
	Pubring *keymgr.Pubring
	// IDDb - Message ID log (replay protection)
	IDDb *idlog.IDLog
	// ChunkDb - Chunk database
	ChunkDb *Chunk
)

func main() {
	var err error
	flag, cfg = config.GetCfg()
	if flag.Version {
		fmt.Println(version)
		os.Exit(0)
	}
	// If the debug flag is set, print the config and exit
	if flag.Debug {
		y, err := cfg.Debug()
		if err != nil {
			fmt.Printf("Debugging Error: %s\n", err)
			os.Exit(1)
		}
		fmt.Printf("%s\n", y)
		os.Exit(0)
	}

	// Set up logging
	loglevel, err := log.Atoi(cfg.General.Loglevel)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: Unknown loglevel", cfg.General.Loglevel)
		os.Exit(1)
	}
	// If we're logging to a file, open the file and redirect output to it
	if cfg.General.LogToFile && cfg.General.LogToJournal {
		fmt.Fprintln(os.Stderr, "Cannot log to file and journal")
		os.Exit(1)
	} else if cfg.General.LogToFile {
		logfile, err := os.OpenFile(cfg.Files.Logfile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%s: Error opening logfile: %v", cfg.Files.Logfile, err)
			os.Exit(1)
		}
		stdlog.SetOutput(logfile)
		log.Current = log.StdLogger{Level: loglevel}
	} else if cfg.General.LogToJournal {
		log.Current = jlog.NewJournal(loglevel)
	} else {
		log.Current = log.StdLogger{Level: loglevel}
	}

	// Inform the user which (if any) config file was used.
	if cfg.Files.Config != "" {
		log.Infof("Using config file: %s", cfg.Files.Config)
	} else {
		log.Warn("No config file was found. Resorting to defaults")
	}

	// Setup complete, time to do some work
	if flag.Client {
		mixprep()
	} else if flag.Stdin {
		dir := maildir.Dir(cfg.Files.Maildir)
		newmsg, err := dir.NewDelivery()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		stdin, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		newmsg.Write(stdin)
		newmsg.Close()
	} else if flag.Remailer {
		err = loopServer()
		if err != nil {
			panic(err)
		}
	} else if flag.Dummy {
		injectDummy()
	} else if flag.Refresh {
		fmt.Printf("Keyring refresh: from=%s, to=%s\n", cfg.Urls.Pubring, cfg.Files.Pubring)
		httpGet(cfg.Urls.Pubring, cfg.Files.Pubring)
		fmt.Printf("Stats refresh: from=%s, to=%s\n", cfg.Urls.Mlist2, cfg.Files.Mlist2)
		httpGet(cfg.Urls.Mlist2, cfg.Files.Mlist2)
	}
	if flag.Send {
		// Flush the outbound pool
		poolOutboundSend()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/go-telegram-bot-api/telegram-bot-api"
)

var ssl = flag.Int("ssl", 0, "0:no ssl, 1:SelfSign, 2:CA-cert")
var webhook = flag.Bool("webhook", false, "webhook mode")
var debug = flag.Bool("debug", false, "debug")
var noisy = flag.Bool("noisy", false, "noisy")
var token = flag.String("token", "", "token")
var pubip = flag.String("pubip", "", "public ip, get with 'curl -s https://ipinfo.io/ip'")
var port = flag.Int("port", 8443, "webhook server port")
var cert = flag.String("cert", "cert.pem", "cert for webhook https server")
var key = flag.String("key", "key.pem", "priv key for webhook https server")
var pathAd = flag.String("Ad", func() string { p, _ := exec.LookPath("Ad"); return p }(), "path to Ad")
var path7z = flag.String("7z", func() string { p, _ := exec.LookPath("7z"); return p }(), "path to 7z")

func handleUpdate(bot *tgbotapi.BotAPI, update tgbotapi.Update) {

	if update.Message == nil {
		return
	}
	Noisy := *noisy
	if update.Message.From.UserName == "sehari24jam" {
		Noisy = true
	}

	log.Printf("[%s] %s", update.Message.From.UserName, update.Message.Text)

	msg := tgbotapi.NewMessage(update.Message.Chat.ID, "Failed")
	msg.ReplyToMessageID = update.Message.MessageID

	switch update.Message.Text {
	case "/start":
		msg.Text = fmt.Sprintf("Welcome %s (%s %s).\n"+
			"You may send me asciidoc (.adoc) file.\n"+
			"Or you can pack whole *.adoc and its included images + sub-adoc into a single compressed file.",
			update.Message.From.UserName, update.Message.From.FirstName, update.Message.From.LastName)
		bot.Send(msg)
		return
	default:
		if update.Message.Document == nil {
			msg.Text = "Send me asciidoc file (.adoc). I don't understand: " + update.Message.Text
			bot.Send(msg)
			return
		}
	}

	f, err := bot.GetFile(tgbotapi.FileConfig{FileID: update.Message.Document.FileID})
	//log.Printf("DocFile: %s", update.Message.Document.FileName)
	if err != nil {
		log.Print(err)
		if Noisy {
			msg.Text = "Failed to proceed uploaded file"
		}
		bot.Send(msg)
		return
	}

	ext := path.Ext(update.Message.Document.FileName)
	tmp := path.Join(os.TempDir(), "ybot."+update.Message.Chat.UserName)
	packed := false
	dpacked := false

	switch strings.ToLower(ext) {

	case ".zip", ".rar", ".7z":
		packed = true
		msg.Text = "Looks good, let me work on this compressed file"

	case ".tgz", ".tbz2", ".txz":
		dpacked = true
		msg.Text = "Looks good, let me work on this compressed file."

	case ".gz", ".bz2", ".xz":
		ftar := strings.TrimSuffix(update.Message.Document.FileName, ext)
		exttar := path.Ext(ftar)
		if exttar == ".tar" {
			dpacked = true
			msg.Text = "Looks good, let me work on this compressed file.."
			ext = exttar + ext
		} else {
			msg.Text = "Document is not an adoc"
			bot.Send(msg)
			return
		}

	case ".adoc":
		msg.Text = "Looks good, let me work on this file"

	default:
		msg.Text = "Document is not an adoc"
		bot.Send(msg)
		return
	}

	if packed || dpacked {
		tmp, err = ioutil.TempDir("", "ybot-")
		if err != nil {
			log.Print(err)
			log.Print(os.RemoveAll(tmp))
			if Noisy {
				msg.Text = "Unable to create temp"
			}
			bot.Send(msg)
			return
		}
	}
	bot.Send(msg)

	workfolder := path.Join(tmp, path.Dir(f.FilePath))
	//lfile := path.Join("/tmp", f.FilePath)
	pdffile := path.Join(workfolder, strings.TrimSuffix(update.Message.Document.FileName, ext)+".pdf")
	workfile := path.Join(workfolder, update.Message.Document.FileName)

	// get WorkFile from TG
	response, err := http.Get(f.Link(*token))
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Failed to get uploaded file"
		}
		bot.Send(msg)
		return
	}
	defer response.Body.Close()

	// create sub folder as necessary
	os.MkdirAll(workfolder, os.ModePerm)

	// save WorkFile
	file, err := os.Create(workfile)
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Unable to create new file"
		}
		bot.Send(msg)
		return
	}
	// Use io.Copy to just dump the response body to the file. This supports huge files
	_, err = io.Copy(file, response.Body)
	file.Close()
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Unable to buffer uploaded file"
		}
		bot.Send(msg)
		return
	}

	go func() {
		///////////////////////////////////// extraction
		switch {
		case packed:
			cmd := exec.Command("7z", "x", workfile)
			cmd.Dir = workfolder
			out, err := cmd.CombinedOutput()
			if err != nil {
				log.Print(err)
				if packed {
					log.Print(os.RemoveAll(tmp))
				}
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v\n%v", string(out), err)
				}
				bot.Send(msg)
				return
			}
			workfile = "*.adoc"
		case dpacked:
			var out bytes.Buffer
			xcmd1 := exec.Command("7z", "x", workfile, "-so")
			xcmd1.Dir = workfolder
			xcmd2 := exec.Command("7z", "x", "-si", "-ttar", "-y")
			xcmd2.Dir = workfolder
			xcmd2.Stdin, _ = xcmd1.StdoutPipe()
			xcmd2.Stdout = &out
			_ = xcmd2.Start()
			_ = xcmd1.Run()
			err := xcmd2.Wait()
			if err != nil {
				log.Print(err)
				if packed {
					log.Print(os.RemoveAll(tmp))
				}
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v\n%v", string(out.Bytes()), err)
				}
				bot.Send(msg)
				return
			}
			workfile = "*.adoc"
		}

		///////////////////////////////////// conversion
		cmd := exec.Command("Ad", workfile)
		cmd.Dir = workfolder
		out, err := cmd.CombinedOutput()
		if err != nil {
			log.Print(err)
			if packed {
				log.Print(os.RemoveAll(tmp))
			}
			if Noisy {
				msg.Text = fmt.Sprintf("Failed %v\n%v", string(out), err)
			}
			bot.Send(msg)
			return
		}

		if packed || dpacked {
			files, err := filepath.Glob(path.Join(workfolder, "*.pdf"))
			if err != nil {
				log.Print(err)
				log.Print(os.RemoveAll(tmp))
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v..", err)
				}
				bot.Send(msg)
				return
			}

			for _, f := range files {
				bot.Send(tgbotapi.NewDocumentUpload(msg.ChatID, f))
			}
			if Noisy {
				msg.Text = fmt.Sprintf("Success %v..", string(out))
			} else {
				msg.Text = "Success.."
			}
			bot.Send(msg)
			log.Print(os.RemoveAll(tmp))
		} else {
			if Noisy {
				msg.Text = fmt.Sprintf("Success %v", string(out))
			} else {
				msg.Text = "Success"
			}
			bot.Send(msg)
			bot.Send(tgbotapi.NewDocumentUpload(msg.ChatID, pdffile))
		}

	}()

}

func main() {

	flag.Parse()

	if *token == "" {
		*token = os.Getenv("YBOTTOKEN")
	}
	bot, err := tgbotapi.NewBotAPI(*token)
	if err != nil {
		log.Panic(err)
	}

	bot.Debug = *debug

	log.Printf("Authorized on account %s", bot.Self.UserName)

	if *webhook {

		var updates tgbotapi.UpdatesChannel
		switch *ssl {
		case 1:
			url := fmt.Sprintf("https://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhookWithCert(url, *cert))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServeTLS(fmt.Sprintf("0.0.0.0:%d", *port), *cert, *key, nil)

		case 2:
			url := fmt.Sprintf("https://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhook(url))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServeTLS(fmt.Sprintf("0.0.0.0:%d", *port), *cert, *key, nil)

		default:
			url := fmt.Sprintf("http://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhook(url))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServe(fmt.Sprintf("0.0.0.0:%d", *port), nil)
		}

		log.Printf("Starting Collect Update from WebHook")
		for update := range updates {
			handleUpdate(bot, update)
			//log.Printf("%+v\n", update)
		}

	} else {

		u := tgbotapi.NewUpdate(0)
		u.Timeout = 60

		updates, err := bot.GetUpdatesChan(u)

		if err != nil {
			log.Panic(err)
		}

		log.Printf("Starting GetUpdate")
		for update := range updates {
			handleUpdate(bot, update)
		}
	}

}
<commit_msg>minor code move<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/go-telegram-bot-api/telegram-bot-api"
)

var ssl = flag.Int("ssl", 0, "0:no ssl, 1:SelfSign, 2:CA-cert")
var webhook = flag.Bool("webhook", false, "webhook mode")
var debug = flag.Bool("debug", false, "debug")
var noisy = flag.Bool("noisy", false, "noisy")
var token = flag.String("token", os.Getenv("YBOTTOKEN"), "token")
var pubip = flag.String("pubip", "", "public ip, get with 'curl -s https://ipinfo.io/ip'")
var port = flag.Int("port", 8443, "webhook server port")
var cert = flag.String("cert", "cert.pem", "cert for webhook https server")
var key = flag.String("key", "key.pem", "priv key for webhook https server")
var pathAd = flag.String("Ad", func() string { p, _ := exec.LookPath("Ad"); return p }(), "path to Ad")
var path7z = flag.String("7z", func() string { p, _ := exec.LookPath("7z"); return p }(), "path to 7z")

func handleUpdate(bot *tgbotapi.BotAPI, update tgbotapi.Update) {

	if update.Message == nil {
		return
	}
	Noisy := *noisy
	if update.Message.From.UserName == "sehari24jam" {
		Noisy = true
	}

	log.Printf("[%s] %s", update.Message.From.UserName, update.Message.Text)

	msg := tgbotapi.NewMessage(update.Message.Chat.ID, "Failed")
	msg.ReplyToMessageID = update.Message.MessageID

	switch update.Message.Text {
	case "/start":
		msg.Text = fmt.Sprintf("Welcome %s (%s %s).\n"+
			"You may send me asciidoc (.adoc) file.\n"+
			"Or you can pack whole *.adoc and its included images + sub-adoc into a single compressed file.",
			update.Message.From.UserName, update.Message.From.FirstName, update.Message.From.LastName)
		bot.Send(msg)
		return
	default:
		if update.Message.Document == nil {
			msg.Text = "Send me asciidoc file (.adoc). I don't understand: " + update.Message.Text
			bot.Send(msg)
			return
		}
	}

	f, err := bot.GetFile(tgbotapi.FileConfig{FileID: update.Message.Document.FileID})
	//log.Printf("DocFile: %s", update.Message.Document.FileName)
	if err != nil {
		log.Print(err)
		if Noisy {
			msg.Text = "Failed to proceed uploaded file"
		}
		bot.Send(msg)
		return
	}

	ext := path.Ext(update.Message.Document.FileName)
	tmp := path.Join(os.TempDir(), "ybot."+update.Message.Chat.UserName)
	packed := false
	dpacked := false

	switch strings.ToLower(ext) {

	case ".zip", ".rar", ".7z":
		packed = true
		msg.Text = "Looks good, let me work on this compressed file"

	case ".tgz", ".tbz2", ".txz":
		dpacked = true
		msg.Text = "Looks good, let me work on this compressed file."

	case ".gz", ".bz2", ".xz":
		ftar := strings.TrimSuffix(update.Message.Document.FileName, ext)
		exttar := path.Ext(ftar)
		if exttar == ".tar" {
			dpacked = true
			msg.Text = "Looks good, let me work on this compressed file.."
			ext = exttar + ext
		} else {
			msg.Text = "Document is not an adoc"
			bot.Send(msg)
			return
		}

	case ".adoc":
		msg.Text = "Looks good, let me work on this file"

	default:
		msg.Text = "Document is not an adoc"
		bot.Send(msg)
		return
	}

	if packed || dpacked {
		tmp, err = ioutil.TempDir("", "ybot-")
		if err != nil {
			log.Print(err)
			log.Print(os.RemoveAll(tmp))
			if Noisy {
				msg.Text = "Unable to create temp"
			}
			bot.Send(msg)
			return
		}
	}
	bot.Send(msg)

	workfolder := path.Join(tmp, path.Dir(f.FilePath))
	//lfile := path.Join("/tmp", f.FilePath)
	pdffile := path.Join(workfolder, strings.TrimSuffix(update.Message.Document.FileName, ext)+".pdf")
	workfile := path.Join(workfolder, update.Message.Document.FileName)

	// get WorkFile from TG
	response, err := http.Get(f.Link(*token))
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Failed to get uploaded file"
		}
		bot.Send(msg)
		return
	}
	defer response.Body.Close()

	// create sub folder as necessary
	os.MkdirAll(workfolder, os.ModePerm)

	// save WorkFile
	file, err := os.Create(workfile)
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Unable to create new file"
		}
		bot.Send(msg)
		return
	}
	// Use io.Copy to just dump the response body to the file. This supports huge files
	_, err = io.Copy(file, response.Body)
	file.Close()
	if err != nil {
		log.Print(err)
		if packed || dpacked {
			log.Print(os.RemoveAll(tmp))
		}
		if Noisy {
			msg.Text = "Unable to buffer uploaded file"
		}
		bot.Send(msg)
		return
	}

	go func() {
		///////////////////////////////////// extraction
		switch {
		case packed:
			cmd := exec.Command("7z", "x", workfile)
			cmd.Dir = workfolder
			out, err := cmd.CombinedOutput()
			if err != nil {
				log.Print(err)
				if packed {
					log.Print(os.RemoveAll(tmp))
				}
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v\n%v", string(out), err)
				}
				bot.Send(msg)
				return
			}
			workfile = "*.adoc"
		case dpacked:
			var out bytes.Buffer
			xcmd1 := exec.Command("7z", "x", "-so", workfile)
			xcmd1.Dir = workfolder
			xcmd2 := exec.Command("7z", "x", "-si", "-ttar", "-y")
			xcmd2.Dir = workfolder
			xcmd2.Stdin, _ = xcmd1.StdoutPipe()
			xcmd2.Stdout = &out
			_ = xcmd2.Start()
			_ = xcmd1.Run()
			err := xcmd2.Wait()
			if err != nil {
				log.Print(err)
				if packed {
					log.Print(os.RemoveAll(tmp))
				}
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v\n%v", string(out.Bytes()), err)
				}
				bot.Send(msg)
				return
			}
			workfile = "*.adoc"
		}

		///////////////////////////////////// conversion
		cmd := exec.Command("Ad", workfile)
		cmd.Dir = workfolder
		out, err := cmd.CombinedOutput()
		if err != nil {
			log.Print(err)
			if packed {
				log.Print(os.RemoveAll(tmp))
			}
			if Noisy {
				msg.Text = fmt.Sprintf("Failed %v\n%v", string(out), err)
			}
			bot.Send(msg)
			return
		}

		if packed || dpacked {
			files, err := filepath.Glob(path.Join(workfolder, "*.pdf"))
			if err != nil {
				log.Print(err)
				log.Print(os.RemoveAll(tmp))
				if Noisy {
					msg.Text = fmt.Sprintf("Failed %v..", err)
				}
				bot.Send(msg)
				return
			}

			for _, f := range files {
				bot.Send(tgbotapi.NewDocumentUpload(msg.ChatID, f))
			}
			if Noisy {
				msg.Text = fmt.Sprintf("Success %v..", string(out))
			} else {
				msg.Text = "Success.."
			}
			bot.Send(msg)
			log.Print(os.RemoveAll(tmp))
		} else {
			if Noisy {
				msg.Text = fmt.Sprintf("Success %v", string(out))
			} else {
				msg.Text = "Success"
			}
			bot.Send(msg)
			bot.Send(tgbotapi.NewDocumentUpload(msg.ChatID, pdffile))
		}

	}()

}

func main() {

	flag.Parse()

	//if *token == "" {
	//	*token = os.Getenv("YBOTTOKEN")
	//}
	bot, err := tgbotapi.NewBotAPI(*token)
	if err != nil {
		log.Panic(err)
	}

	bot.Debug = *debug

	log.Printf("Authorized on account %s", bot.Self.UserName)

	if *webhook {

		var updates tgbotapi.UpdatesChannel
		switch *ssl {
		case 1:
			url := fmt.Sprintf("https://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhookWithCert(url, *cert))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServeTLS(fmt.Sprintf("0.0.0.0:%d", *port), *cert, *key, nil)

		case 2:
			url := fmt.Sprintf("https://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhook(url))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServeTLS(fmt.Sprintf("0.0.0.0:%d", *port), *cert, *key, nil)

		default:
			url := fmt.Sprintf("http://%s:%d/%s", *pubip, *port, bot.Token)
			log.Print("Webhook URL: " + url)
			_, err = bot.SetWebhook(tgbotapi.NewWebhook(url))
			if err != nil {
				log.Fatal(err)
			}

			updates = bot.ListenForWebhook("/" + bot.Token)
			go http.ListenAndServe(fmt.Sprintf("0.0.0.0:%d", *port), nil)
		}

		log.Printf("Starting Collect Update from WebHook")
		for update := range updates {
			handleUpdate(bot, update)
			//log.Printf("%+v\n", update)
		}

	} else {

		u := tgbotapi.NewUpdate(0)
		u.Timeout = 60

		updates, err := bot.GetUpdatesChan(u)

		if err != nil {
			log.Panic(err)
		}

		log.Printf("Starting GetUpdate")
		for update := range updates {
			handleUpdate(bot, update)
		}
	}

}
<|endoftext|>


<commit_before>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The seeddb command populates a database with an initial set of modules.
package main

import (
	"context"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"

	_ "github.com/jackc/pgx/v4/stdlib" // for pgx driver
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/config/dynconfig"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/stdlib"
	"golang.org/x/pkgsite/internal/worker"
	"golang.org/x/sync/errgroup"
)

var (
	seedfile = flag.String("seed", "devtools/cmd/seeddb/seed.txt", "filename containing modules for seeding the database")
	refetch  = flag.Bool("refetch", false, "refetch modules in the seedfile even if they already exist")
)

func main() {
	flag.Parse()

	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}

	exps, err := fetchExperiments(ctx, cfg)
	if err != nil {
		log.Fatal(ctx, err)
	}
	ctx = experiment.NewContext(ctx, exps...)

	db, err := database.Open("pgx", cfg.DBConnInfo(), "seeddb")
	if err != nil {
		log.Fatalf(ctx, "database.Open for host %s failed with %v", cfg.DBHost, err)
	}
	defer db.Close()

	if err := run(ctx, db, cfg.ProxyURL); err != nil {
		log.Fatal(ctx, err)
	}
}

func run(ctx context.Context, db *database.DB, proxyURL string) error {
	start := time.Now()

	proxyClient, err := proxy.New(proxyURL)
	if err != nil {
		return err
	}

	sourceClient := source.NewClient(config.SourceTimeout)
	seedModules, err := readSeedFile(ctx, *seedfile)
	if err != nil {
		return err
	}

	// Expand versions and group by module path.
	versionsByPath := map[string][]string{}
	for _, m := range seedModules {
		vers, err := versions(ctx, proxyClient, m)
		if err != nil {
			return err
		}
		versionsByPath[m.Path] = append(versionsByPath[m.Path], vers...)
	}

	r := results{}
	g := new(errgroup.Group)
	f := &worker.Fetcher{
		ProxyClient:  proxyClient,
		SourceClient: sourceClient,
		DB:           postgres.New(db),
	}
	for path, vers := range versionsByPath {
		path := path
		vers := vers
		// Process versions of the same module sequentially, to avoid DB contention.
		g.Go(func() error {
			for _, v := range vers {
				if err := fetch(ctx, db, f, path, v, &r); err != nil {
					return err
				}
			}
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	log.Infof(ctx, "Successfully fetched all modules: %v", time.Since(start))

	// Print the time it took to fetch these modules.
	var keys []string
	for k := range r.paths {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		log.Infof(ctx, "%s | %v", k, r.paths[k])
	}
	return nil
}

func versions(ctx context.Context, proxyClient *proxy.Client, mv internal.Modver) ([]string, error) {
	if mv.Version != "all" {
		return []string{mv.Version}, nil
	}
	if mv.Path == stdlib.ModulePath {
		stdVersions, err := stdlib.Versions()
		if err != nil {
			return nil, err
		}
		// As an optimization, only fetch release versions for the standard
		// library.
		var vers []string
		for _, v := range stdVersions {
			if strings.HasSuffix(v, ".0") {
				vers = append(vers, v)
			}
		}
		return vers, nil
	}
	return proxyClient.Versions(ctx, mv.Path)
}

func fetch(ctx context.Context, db *database.DB, f *worker.Fetcher, m, v string, r *results) error {
	// Record the duration of this fetch request.
	start := time.Now()

	var exists bool
	defer func() {
		r.add(m, v, start, exists)
	}()
	err := db.QueryRow(ctx, `
		SELECT 1 FROM modules WHERE module_path = $1 AND version = $2;
	`, m, v).Scan(&exists)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return err
	}
	if errors.Is(err, sql.ErrNoRows) || *refetch {
		return fetchFunc(ctx, f, m, v)
	}
	return nil
}

func fetchFunc(ctx context.Context, f *worker.Fetcher, m, v string) (err error) {
	defer derrors.Wrap(&err, "fetchFunc(ctx, f, %q, %q)", m, v)

	log.Infof(ctx, "Fetch requested: %q %q", m, v)
	fetchCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	code, _, err := f.FetchAndUpdateState(fetchCtx, m, v, "")
	if err != nil {
		if code == http.StatusNotFound {
			// We expect
			// github.com/jackc/pgx/pgxpool@v3.6.2+incompatible
			// to fail from seed.txt, so that it will redirect to
			// github.com/jackc/pgx/v4/pgxpool in tests.
			return nil
		}
		return err
	}
	return nil
}

type results struct {
	mu    sync.Mutex
	paths map[string]time.Duration
}

func (r *results) add(modPath, version string, start time.Time, exists bool) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.paths == nil {
		r.paths = map[string]time.Duration{}
	}
	key := fmt.Sprintf("%s@%s", modPath, version)
	if exists {
		key = fmt.Sprintf("%s (exists)", key)
	}
	r.paths[key] = time.Since(start)
}

// readSeedFile reads a file of module versions that we want to fetch for
// seeding the database. Each line of the file should be of the form:
//     module@version
func readSeedFile(ctx context.Context, seedfile string) (_ []internal.Modver, err error) {
	defer derrors.Wrap(&err, "readSeedFile %q", seedfile)
	lines, err := internal.ReadFileLines(seedfile)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "read %d module versions from %s", len(lines), seedfile)

	var modules []internal.Modver
	for _, l := range lines {
		mv, err := internal.ParseModver(l)
		if err != nil {
			return nil, err
		}
		modules = append(modules, mv)
	}
	return modules, nil
}

func fetchExperiments(ctx context.Context, cfg *config.Config) ([]string, error) {
	if cfg.DynamicConfigLocation == "" {
		return nil, nil
	}
	dc, err := dynconfig.Read(ctx, cfg.DynamicConfigLocation)
	if err != nil {
		return nil, err
	}
	var exps []string
	for _, e := range dc.Experiments {
		if e.Rollout > 0 {
			exps = append(exps, e.Name)
		}
	}
	return exps, nil
}
<commit_msg>devtools/cmd/seeddb: add -keep_going<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The seeddb command populates a database with an initial set of modules.
package main

import (
	"context"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"

	_ "github.com/jackc/pgx/v4/stdlib" // for pgx driver
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/config/dynconfig"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/stdlib"
	"golang.org/x/pkgsite/internal/worker"
	"golang.org/x/sync/errgroup"
)

var (
	seedfile  = flag.String("seed", "devtools/cmd/seeddb/seed.txt", "filename containing modules for seeding the database")
	refetch   = flag.Bool("refetch", false, "refetch modules in the seedfile even if they already exist")
	keepGoing = flag.Bool("keep_going", false, "continue on errors")
)

func main() {
	flag.Parse()

	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}

	exps, err := fetchExperiments(ctx, cfg)
	if err != nil {
		log.Fatal(ctx, err)
	}
	ctx = experiment.NewContext(ctx, exps...)

	db, err := database.Open("pgx", cfg.DBConnInfo(), "seeddb")
	if err != nil {
		log.Fatalf(ctx, "database.Open for host %s failed with %v", cfg.DBHost, err)
	}
	defer db.Close()

	if err := run(ctx, db, cfg.ProxyURL); err != nil {
		log.Fatal(ctx, err)
	}
}

func run(ctx context.Context, db *database.DB, proxyURL string) error {
	start := time.Now()

	proxyClient, err := proxy.New(proxyURL)
	if err != nil {
		return err
	}

	sourceClient := source.NewClient(config.SourceTimeout)
	seedModules, err := readSeedFile(ctx, *seedfile)
	if err != nil {
		return err
	}

	// Expand versions and group by module path.
	versionsByPath := map[string][]string{}
	for _, m := range seedModules {
		vers, err := versions(ctx, proxyClient, m)
		if err != nil {
			return err
		}
		versionsByPath[m.Path] = append(versionsByPath[m.Path], vers...)
	}

	r := results{}
	g := new(errgroup.Group)
	f := &worker.Fetcher{
		ProxyClient:  proxyClient,
		SourceClient: sourceClient,
		DB:           postgres.New(db),
	}
	var (
		mu     sync.Mutex
		errors database.MultiErr
	)
	for path, vers := range versionsByPath {
		path := path
		vers := vers
		// Process versions of the same module sequentially, to avoid DB contention.
		g.Go(func() error {
			for _, v := range vers {
				if err := fetch(ctx, db, f, path, v, &r); err != nil {
					if *keepGoing {
						mu.Lock()
						errors = append(errors, err)
						mu.Unlock()
					} else {
						return err
					}
				}
			}
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	if len(errors) > 0 {
		return errors
	}
	log.Infof(ctx, "Successfully fetched all modules: %v", time.Since(start))

	// Print the time it took to fetch these modules.
	var keys []string
	for k := range r.paths {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		log.Infof(ctx, "%s | %v", k, r.paths[k])
	}
	return nil
}

func versions(ctx context.Context, proxyClient *proxy.Client, mv internal.Modver) ([]string, error) {
	if mv.Version != "all" {
		return []string{mv.Version}, nil
	}
	if mv.Path == stdlib.ModulePath {
		stdVersions, err := stdlib.Versions()
		if err != nil {
			return nil, err
		}
		// As an optimization, only fetch release versions for the standard
		// library.
		var vers []string
		for _, v := range stdVersions {
			if strings.HasSuffix(v, ".0") {
				vers = append(vers, v)
			}
		}
		return vers, nil
	}
	return proxyClient.Versions(ctx, mv.Path)
}

func fetch(ctx context.Context, db *database.DB, f *worker.Fetcher, m, v string, r *results) error {
	// Record the duration of this fetch request.
	start := time.Now()

	var exists bool
	defer func() {
		r.add(m, v, start, exists)
	}()
	err := db.QueryRow(ctx, `
		SELECT 1 FROM modules WHERE module_path = $1 AND version = $2;
	`, m, v).Scan(&exists)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return err
	}
	if errors.Is(err, sql.ErrNoRows) || *refetch {
		return fetchFunc(ctx, f, m, v)
	}
	return nil
}

func fetchFunc(ctx context.Context, f *worker.Fetcher, m, v string) (err error) {
	defer derrors.Wrap(&err, "fetchFunc(ctx, f, %q, %q)", m, v)

	log.Infof(ctx, "Fetch requested: %q %q", m, v)
	fetchCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	code, _, err := f.FetchAndUpdateState(fetchCtx, m, v, "")
	if err != nil {
		if code == http.StatusNotFound {
			// We expect
			// github.com/jackc/pgx/pgxpool@v3.6.2+incompatible
			// to fail from seed.txt, so that it will redirect to
			// github.com/jackc/pgx/v4/pgxpool in tests.
			return nil
		}
		return err
	}
	return nil
}

type results struct {
	mu    sync.Mutex
	paths map[string]time.Duration
}

func (r *results) add(modPath, version string, start time.Time, exists bool) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.paths == nil {
		r.paths = map[string]time.Duration{}
	}
	key := fmt.Sprintf("%s@%s", modPath, version)
	if exists {
		key = fmt.Sprintf("%s (exists)", key)
	}
	r.paths[key] = time.Since(start)
}

// readSeedFile reads a file of module versions that we want to fetch for
// seeding the database. Each line of the file should be of the form:
//     module@version
func readSeedFile(ctx context.Context, seedfile string) (_ []internal.Modver, err error) {
	defer derrors.Wrap(&err, "readSeedFile %q", seedfile)
	lines, err := internal.ReadFileLines(seedfile)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "read %d module versions from %s", len(lines), seedfile)

	var modules []internal.Modver
	for _, l := range lines {
		mv, err := internal.ParseModver(l)
		if err != nil {
			return nil, err
		}
		modules = append(modules, mv)
	}
	return modules, nil
}

func fetchExperiments(ctx context.Context, cfg *config.Config) ([]string, error) {
	if cfg.DynamicConfigLocation == "" {
		return nil, nil
	}
	dc, err := dynconfig.Read(ctx, cfg.DynamicConfigLocation)
	if err != nil {
		return nil, err
	}
	var exps []string
	for _, e := range dc.Experiments {
		if e.Rollout > 0 {
			exps = append(exps, e.Name)
		}
	}
	return exps, nil
}
<|endoftext|>


<commit_before>// web100lib provides Go bindings to some functions in the web100 library.
package web100lib

// Cgo directives must immediately preceed 'import "C"' below.
// For more information see:
//  - https://blog.golang.org/c-go-cgo
//  - https://golang.org/cmd/cgo

/*
#include <stdio.h>
#include <stdlib.h>
#include <web100.h>
#include <web100-int.h>

#include <arpa/inet.h>
*/
import "C"

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"unsafe"

	"cloud.google.com/go/bigquery"
)

// Discoveries:
//  - Not all C macros exist in the "C" namespace.
//  - 'NULL' is usually equivalent to 'nil'

// Web100 maintains state associated with a web100 log file.
type Web100 struct {
	// legacyNames maps legacy web100 variable names to their canonical names.
	legacyNames map[string]string

	// Do not export unsafe pointers.
	log  unsafe.Pointer
	snap unsafe.Pointer
}

// Open prepares a web100 log file for reading. The caller must call Close on
// the returned Web100 instance to release resources.
func Open(filename string, legacyNames map[string]string) (*Web100, error) {
	c_filename := C.CString(filename)
	defer C.free(unsafe.Pointer(c_filename))

	log := C.web100_log_open_read(c_filename)
	if log == nil {
		return nil, fmt.Errorf(C.GoString(C.web100_strerror(C.web100_errno)))
	}

	// Pre-allocate a snapshot record.
	snap := C.web100_snapshot_alloc_from_log(log)

	w := &Web100{
		legacyNames: legacyNames,
		log:         unsafe.Pointer(log),
		snap:        unsafe.Pointer(snap),
	}
	return w, nil
}

// Next iterates through the web100 log file reading the next snapshot record
// until EOF or an error occurs.
func (w *Web100) Next() error {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	// Read the next web100_snaplog data from underlying file.
	err := C.web100_snap_from_log(snap, log)
	if err == C.EOF {
		return io.EOF
	}
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}
	return nil
}

func (w *Web100) Values() (map[string]bigquery.Value, error) {
	v, err := w.logValues()
	if err != nil {
		return nil, err
	}
	v, err = w.snapValues(v)
	if err != nil {
		return nil, err
	}
	return v, nil
}

// logValues returns a map of values from the web100 log. IPv6 address
// connection information is not available and must be set based on a snapshot.
func (w *Web100) logValues() (map[string]bigquery.Value, error) {
	log := (*C.web100_log)(w.log)

	agent := C.web100_get_log_agent(log)

	results := make(map[string]bigquery.Value)
	results["web100_log_entry.version"] = C.GoString(C.web100_get_agent_version(agent))

	time := C.web100_get_log_time(log)
	results["web100_log_entry.log_time"] = int64(time)

	conn := C.web100_get_log_connection(log)
	// NOTE: web100_connection_spec_v6 is not filled in by the web100 library.
	// NOTE: addrtype is always WEB100_ADDRTYPE_UNKNOWN.
	// NOTE: legacy values for local_af are: IPv4 = 0, IPv6 = 1.
	results["web100_log_entry.connection_spec.local_af"] = int64(0)

	var spec C.struct_web100_connection_spec
	C.web100_get_connection_spec(conn, &spec)

	addr := C.struct_in_addr{C.in_addr_t(spec.src_addr)}
	results["web100_log_entry.connection_spec.local_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.local_port"] = int64(spec.src_port)

	addr = C.struct_in_addr{C.in_addr_t(spec.dst_addr)}
	results["web100_log_entry.connection_spec.remote_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.remote_port"] = int64(spec.dst_port)

	return results, nil
}

// snapValues converts all variables in the latest snap record into a results map.
func (w *Web100) snapValues(logValues map[string]bigquery.Value) (map[string]bigquery.Value, error) {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	var_text := C.calloc(2*C.WEB100_VALUE_LEN_MAX, 1) // Use a better size.
	defer C.free(var_text)

	var_data := C.calloc(C.WEB100_VALUE_LEN_MAX, 1)
	defer C.free(var_data)

	// Parses variables from most recent web100_snapshot data.
	group := C.web100_get_log_group(log)
	for v := C.web100_var_head(group); v != nil; v = C.web100_var_next(v) {

		name := C.web100_get_var_name(v)
		var_type := C.web100_get_var_type(v)

		// Read the raw variable data from the snapshot data.
		errno := C.web100_snap_read(v, snap, var_data)
		if errno != C.WEB100_ERR_SUCCESS {
			return nil, fmt.Errorf(C.GoString(C.web100_strerror(errno)))
		}

		// Convert raw var_data into a string based on var_type.
		// TODO(final): ultimately, we should reimplement web100_value_to_textn to operate on Go types.
		C.web100_value_to_textn((*C.char)(var_text), C.WEB100_VALUE_LEN_MAX, (C.WEB100_TYPE)(var_type), var_data)

		// Use the canonical variable name.
		var canonicalName string
		if _, ok := w.legacyNames[canonicalName]; ok {
			canonicalName = w.legacyNames[canonicalName]
		} else {
			canonicalName = C.GoString(name)
		}

		// Attempt to convert the current variable to an int64.
		value, err := strconv.ParseInt(C.GoString((*C.char)(var_text)), 10, 64)
		if err != nil {
			// Leave variable as a string.
			fmt.Println("var_type:", int(var_type), "name:", canonicalName)
			logValues[fmt.Sprintf("web100_log_entry.snap.%s", canonicalName)] = C.GoString((*C.char)(var_text))
		} else {
			//
			logValues[fmt.Sprintf("web100_log_entry.snap.%s", canonicalName)] = value
		}
	}
	return logValues, nil
}

// Close releases resources created by Open.
func (w *Web100) Close() error {
	snap := (*C.web100_snapshot)(w.snap)
	C.web100_snapshot_free(snap)

	log := (*C.web100_log)(w.log)
	err := C.web100_log_close_read(log)
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}

	// Clear pointer after free.
	w.log = nil
	w.snap = nil
	return nil
}

func LookupError(errnum int) string {
	return C.GoString(C.web100_strerror(C.int(errnum)))
}

func PrettyPrint(results map[string]string) {
	b, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	fmt.Print(string(b))
}
<commit_msg>Remove log line.<commit_after>// web100lib provides Go bindings to some functions in the web100 library.
package web100lib

// Cgo directives must immediately preceed 'import "C"' below.
// For more information see:
//  - https://blog.golang.org/c-go-cgo
//  - https://golang.org/cmd/cgo

/*
#include <stdio.h>
#include <stdlib.h>
#include <web100.h>
#include <web100-int.h>

#include <arpa/inet.h>
*/
import "C"

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"unsafe"

	"cloud.google.com/go/bigquery"
)

// Discoveries:
//  - Not all C macros exist in the "C" namespace.
//  - 'NULL' is usually equivalent to 'nil'

// Web100 maintains state associated with a web100 log file.
type Web100 struct {
	// legacyNames maps legacy web100 variable names to their canonical names.
	legacyNames map[string]string

	// Do not export unsafe pointers.
	log  unsafe.Pointer
	snap unsafe.Pointer
}

// Open prepares a web100 log file for reading. The caller must call Close on
// the returned Web100 instance to release resources.
func Open(filename string, legacyNames map[string]string) (*Web100, error) {
	c_filename := C.CString(filename)
	defer C.free(unsafe.Pointer(c_filename))

	log := C.web100_log_open_read(c_filename)
	if log == nil {
		return nil, fmt.Errorf(C.GoString(C.web100_strerror(C.web100_errno)))
	}

	// Pre-allocate a snapshot record.
	snap := C.web100_snapshot_alloc_from_log(log)

	w := &Web100{
		legacyNames: legacyNames,
		log:         unsafe.Pointer(log),
		snap:        unsafe.Pointer(snap),
	}
	return w, nil
}

// Next iterates through the web100 log file reading the next snapshot record
// until EOF or an error occurs.
func (w *Web100) Next() error {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	// Read the next web100_snaplog data from underlying file.
	err := C.web100_snap_from_log(snap, log)
	if err == C.EOF {
		return io.EOF
	}
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}
	return nil
}

func (w *Web100) Values() (map[string]bigquery.Value, error) {
	v, err := w.logValues()
	if err != nil {
		return nil, err
	}
	v, err = w.snapValues(v)
	if err != nil {
		return nil, err
	}
	return v, nil
}

// logValues returns a map of values from the web100 log. IPv6 address
// connection information is not available and must be set based on a snapshot.
func (w *Web100) logValues() (map[string]bigquery.Value, error) {
	log := (*C.web100_log)(w.log)

	agent := C.web100_get_log_agent(log)

	results := make(map[string]bigquery.Value)
	results["web100_log_entry.version"] = C.GoString(C.web100_get_agent_version(agent))

	time := C.web100_get_log_time(log)
	results["web100_log_entry.log_time"] = int64(time)

	conn := C.web100_get_log_connection(log)
	// NOTE: web100_connection_spec_v6 is not filled in by the web100 library.
	// NOTE: addrtype is always WEB100_ADDRTYPE_UNKNOWN.
	// NOTE: legacy values for local_af are: IPv4 = 0, IPv6 = 1.
	results["web100_log_entry.connection_spec.local_af"] = int64(0)

	var spec C.struct_web100_connection_spec
	C.web100_get_connection_spec(conn, &spec)

	addr := C.struct_in_addr{C.in_addr_t(spec.src_addr)}
	results["web100_log_entry.connection_spec.local_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.local_port"] = int64(spec.src_port)

	addr = C.struct_in_addr{C.in_addr_t(spec.dst_addr)}
	results["web100_log_entry.connection_spec.remote_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.remote_port"] = int64(spec.dst_port)

	return results, nil
}

// snapValues converts all variables in the latest snap record into a results map.
func (w *Web100) snapValues(logValues map[string]bigquery.Value) (map[string]bigquery.Value, error) {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	var_text := C.calloc(2*C.WEB100_VALUE_LEN_MAX, 1) // Use a better size.
	defer C.free(var_text)

	var_data := C.calloc(C.WEB100_VALUE_LEN_MAX, 1)
	defer C.free(var_data)

	// Parses variables from most recent web100_snapshot data.
	group := C.web100_get_log_group(log)
	for v := C.web100_var_head(group); v != nil; v = C.web100_var_next(v) {

		name := C.web100_get_var_name(v)
		var_type := C.web100_get_var_type(v)

		// Read the raw variable data from the snapshot data.
		errno := C.web100_snap_read(v, snap, var_data)
		if errno != C.WEB100_ERR_SUCCESS {
			return nil, fmt.Errorf(C.GoString(C.web100_strerror(errno)))
		}

		// Convert raw var_data into a string based on var_type.
		// TODO(final): ultimately, we should reimplement web100_value_to_textn to operate on Go types.
		C.web100_value_to_textn((*C.char)(var_text), C.WEB100_VALUE_LEN_MAX, (C.WEB100_TYPE)(var_type), var_data)

		// Use the canonical variable name.
		var canonicalName string
		if _, ok := w.legacyNames[canonicalName]; ok {
			canonicalName = w.legacyNames[canonicalName]
		} else {
			canonicalName = C.GoString(name)
		}

		// Attempt to convert the current variable to an int64.
		value, err := strconv.ParseInt(C.GoString((*C.char)(var_text)), 10, 64)
		if err != nil {
			// Leave variable as a string.
			logValues[fmt.Sprintf("web100_log_entry.snap.%s", canonicalName)] = C.GoString((*C.char)(var_text))
		} else {
			//
			logValues[fmt.Sprintf("web100_log_entry.snap.%s", canonicalName)] = value
		}
	}
	return logValues, nil
}

// Close releases resources created by Open.
func (w *Web100) Close() error {
	snap := (*C.web100_snapshot)(w.snap)
	C.web100_snapshot_free(snap)

	log := (*C.web100_log)(w.log)
	err := C.web100_log_close_read(log)
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}

	// Clear pointer after free.
	w.log = nil
	w.snap = nil
	return nil
}

func LookupError(errnum int) string {
	return C.GoString(C.web100_strerror(C.int(errnum)))
}

func PrettyPrint(results map[string]string) {
	b, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	fmt.Print(string(b))
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// The format tests are white box tests, meaning that the tests are in the
// same package as the code, as all the format details are internal to the
// package.

package agent

import (
	"os"
	"path"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
)

type format112Suite struct {
	testing.LoggingSuite
	formatter formatter112
}

var _ = gc.Suite(&format112Suite{})

var agentParams = AgentConfigParams{
	Tag:            "omg",
	Password:       "sekrit",
	CACert:         []byte("ca cert"),
	StateAddresses: []string{"localhost:1234"},
	APIAddresses:   []string{"localhost:1235"},
	Nonce:          "a nonce",
}

func (s *format112Suite) newConfig(c *gc.C) *configInternal {
	params := agentParams
	params.DataDir = c.MkDir()
	config, err := newConfig(params)
	c.Assert(err, gc.IsNil)
	return config
}

func (s *format112Suite) TestWriteAgentConfig(c *gc.C) {
	config := s.newConfig(c)
	err := s.formatter.write(config)
	c.Assert(err, gc.IsNil)

	expectedLocation := path.Join(config.Dir(), "agent.conf")
	fileInfo, err := os.Stat(expectedLocation)
	c.Assert(err, gc.IsNil)
	c.Assert(fileInfo.Mode().IsRegular(), jc.IsTrue)
	c.Assert(fileInfo.Mode().Perm(), gc.Equals, os.FileMode(0600))
	c.Assert(fileInfo.Size(), jc.GreaterThan, 0)
}

func (s *format112Suite) TestRead(c *gc.C) {
	config := s.newConfig(c)
	err := s.formatter.write(config)
	c.Assert(err, gc.IsNil)
	// The readConfig is missing the dataDir initially.
	readConfig, err := s.formatter.read(config.Dir())
	c.Assert(err, gc.IsNil)
	c.Assert(readConfig.dataDir, gc.Equals, "")
	// This is put in by the ReadConf method that we are avoiding using
	// becuase it will have side-effects soon around migrating configs.
	readConfig.dataDir = config.dataDir
	c.Assert(readConfig, gc.DeepEquals, config)
}

func (s *format112Suite) TestWriteCommands(c *gc.C) {
	config := s.newConfig(c)
	commands, err := s.formatter.writeCommands(config)
	c.Assert(err, gc.IsNil)
	c.Assert(commands, gc.HasLen, 3)
	c.Assert(commands[0], gc.Matches, `mkdir -p '\S+/agents/omg'`)
	c.Assert(commands[1], gc.Matches, `install -m 600 /dev/null '\S+/agents/omg/agent.conf'`)
	c.Assert(commands[2], gc.Matches, `printf '%s\\n' '(.|\n)*' > '\S+/agents/omg/agent.conf'`)
}
<commit_msg>Test state config.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// The format tests are white box tests, meaning that the tests are in the
// same package as the code, as all the format details are internal to the
// package.

package agent

import (
	"os"
	"path"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
)

type format112Suite struct {
	testing.LoggingSuite
	formatter formatter112
}

var _ = gc.Suite(&format112Suite{})

var agentParams = AgentConfigParams{
	Tag:            "omg",
	Password:       "sekrit",
	CACert:         []byte("ca cert"),
	StateAddresses: []string{"localhost:1234"},
	APIAddresses:   []string{"localhost:1235"},
	Nonce:          "a nonce",
}

func (s *format112Suite) newConfig(c *gc.C) *configInternal {
	params := agentParams
	params.DataDir = c.MkDir()
	config, err := newConfig(params)
	c.Assert(err, gc.IsNil)
	return config
}

func (s *format112Suite) TestWriteAgentConfig(c *gc.C) {
	config := s.newConfig(c)
	err := s.formatter.write(config)
	c.Assert(err, gc.IsNil)

	expectedLocation := path.Join(config.Dir(), "agent.conf")
	fileInfo, err := os.Stat(expectedLocation)
	c.Assert(err, gc.IsNil)
	c.Assert(fileInfo.Mode().IsRegular(), jc.IsTrue)
	c.Assert(fileInfo.Mode().Perm(), gc.Equals, os.FileMode(0600))
	c.Assert(fileInfo.Size(), jc.GreaterThan, 0)
}

func (s *format112Suite) assertWriteAndRead(c *gc.C, config *configInternal) {
	err := s.formatter.write(config)
	c.Assert(err, gc.IsNil)
	// The readConfig is missing the dataDir initially.
	readConfig, err := s.formatter.read(config.Dir())
	c.Assert(err, gc.IsNil)
	c.Assert(readConfig.dataDir, gc.Equals, "")
	// This is put in by the ReadConf method that we are avoiding using
	// becuase it will have side-effects soon around migrating configs.
	readConfig.dataDir = config.dataDir
	c.Assert(readConfig, gc.DeepEquals, config)
}

func (s *format112Suite) TestRead(c *gc.C) {
	config := s.newConfig(c)
	s.assertWriteAndRead(c, config)
}

func (s *format112Suite) TestWriteCommands(c *gc.C) {
	config := s.newConfig(c)
	commands, err := s.formatter.writeCommands(config)
	c.Assert(err, gc.IsNil)
	c.Assert(commands, gc.HasLen, 3)
	c.Assert(commands[0], gc.Matches, `mkdir -p '\S+/agents/omg'`)
	c.Assert(commands[1], gc.Matches, `install -m 600 /dev/null '\S+/agents/omg/agent.conf'`)
	c.Assert(commands[2], gc.Matches, `printf '%s\\n' '(.|\n)*' > '\S+/agents/omg/agent.conf'`)
}

func (s *format112Suite) TestReadWriteStateConfig(c *gc.C) {
	stateParams := StateMachineConfigParams{
		AgentConfigParams: agentParams,
		StateServerCert:   []byte("some special cert"),
		StateServerKey:    []byte("a special key"),
		StatePort:         12345,
		APIPort:           23456,
	}
	stateParams.DataDir = c.MkDir()
	configInterface, err := NewStateMachineConfig(stateParams)
	c.Assert(err, gc.IsNil)
	config, ok := configInterface.(*configInternal)
	c.Assert(ok, jc.IsTrue)

	s.assertWriteAndRead(c, config)
}
<|endoftext|>