<commit_before>package main

import (
	//DB
	"database/sql"

	"github.com/golang/glog"
	"github.com/iambc/xerrors"
	_ "github.com/lib/pq"
)

//Relational database implementation for writer interface
type writerrdb struct {
	*sql.DB
}

func (db *writerrdb) getBoards(apiKey string) (currBoards []boards, err error) {
	glog.Info(" apiKey: ", apiKey)
	rows, err := db.Query("select b.ID, b.name, b.descr from boards b join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID where api_key = $1;", apiKey)
	if err != nil {
		return currBoards, xerrors.NewUIErr(err.Error(), err.Error(), `002`, true)
	}
	defer rows.Close()

	for rows.Next() {
		var board boards
		err = rows.Scan(&board.ID, &board.Name, &board.Descr)
		if err != nil {
			return currBoards, xerrors.NewUIErr(err.Error(), err.Error(), `003`, true)
		}
		currBoards = append(currBoards, board)
	}
	return currBoards, nil
}

func (db *writerrdb) getActiveThreadsForBoard(apiKey string, boardID int) (activeThreads []threads, err error) {
	rows, err := db.Query(`select t.ID, t.name, count(*), (select count(*) from thread_posts where thread_ID = t.ID and attachment_url is not null) from threads t
				join boards b on b.ID = t.board_ID
				join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID
				left join thread_posts tp on tp.thread_ID = t.ID
			    where t.is_active = TRUE and t.board_ID = $1 and ibc.api_key = $2 group by 1,2 order by t.ID;`, boardID, apiKey)
	if err != nil {
		return activeThreads, xerrors.NewUIErr(err.Error(), err.Error(), `006`, true)
	}
	defer rows.Close()

	for rows.Next() {
		var thread threads
		err = rows.Scan(&thread.ID, &thread.Name, &thread.PostCount, &thread.PostCountWithAttachment)
		if err != nil {
			return activeThreads, xerrors.NewUIErr(err.Error(), err.Error(), `007`, true)
		}
		activeThreads = append(activeThreads, thread)
	}
	return activeThreads, nil
}

func (db *writerrdb) getPostsForThread(apiKey string, threadID int) (currPosts []threadPosts, err error) {
	rows, err := db.Query(`select tp.ID, tp.body, tp.attachment_url, tp.inserted_at, tp.source_ip
			    from thread_posts tp join threads t on t.ID = tp.thread_ID
						 join boards b on b.ID = t.board_ID
						 join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID
			    where tp.thread_ID = $1 and ibc.api_key = $2 and t.is_active = true;`, threadID, apiKey)
	if err != nil {
		glog.Error(err)
		return currPosts, xerrors.NewSysErr()
	}
	defer rows.Close()

	for rows.Next() {
		var currPost threadPosts
		err = rows.Scan(&currPost.ID, &currPost.Body, &currPost.attachmentURL, &currPost.InsertedAt, &currPost.SourceIP)
		if err != nil {
			glog.Error(err)
			return currPosts, xerrors.NewSysErr()
		}
		currPosts = append(currPosts, currPost)
	}
	return currPosts, err
}

func (db *writerrdb) addPostToThread(threadID int, threadBodyPost string, attachmentURL *string, clientRemoteAddr string) (err error) {
	_, err = db.Query("INSERT INTO thread_posts(body, thread_ID, attachment_url, source_ip) VALUES($1, $2, $3, $4)", threadBodyPost, threadID, attachmentURL, clientRemoteAddr)

	if err != nil {
		glog.Error(err)
		return xerrors.NewUIErr(err.Error(), err.Error(), `011`, true)
	}
	return nil
}

func (db *writerrdb) addThread(boardID int, threadName string) (threads, error) {

	var threadID int
	err := db.QueryRow("INSERT INTO threads(name, board_ID, limits_reached_action_ID, max_posts_per_thread) VALUES($1, $2, 1, 10)  RETURNING ID, name", threadName, boardID).Scan(&threadID, &threadName)

	if err != nil {
		glog.Error("INSERT FAILED")
		return threads{ID: -1, Name: `err`}, xerrors.NewUIErr(err.Error(), err.Error(), `017`, true)
	}
	return threads{ID: threadID, Name: threadName}, nil
}

func (db *writerrdb) isThreadLimitReached(boardID int) (bool, error) {
	var isLimitReached bool
	err := db.QueryRow("select (select count(*) from threads  where board_ID = $1) > thread_setting_max_thread_count  from boards where ID = $1;", boardID).Scan(&isLimitReached)
	if err != nil {
		glog.Error("COULD NOT SELECT thread_count")
		return true, xerrors.NewUIErr(err.Error(), err.Error(), `015`, true)
	}

	return isLimitReached, nil
}

func (db *writerrdb) isPostLimitReached(threadID int) (bool, threads, error) {
	var isLimitReached bool
	var thread threads
	err := db.QueryRow("select (select count(*) from thread_posts  where thread_ID = $1) > max_posts_per_thread, min_post_length, max_post_length  from threads where ID = $1;", threadID).Scan(&isLimitReached, &thread.MinPostLength, &thread.MaxPostLength)
	if err != nil {
		return true, thread, xerrors.NewUIErr(err.Error(), err.Error(), `009`, true)
	}
	return isLimitReached, thread, err
}

func (db *writerrdb) archiveThread(threadID int) {
	db.QueryRow("UPDATE threads set is_active = false where ID = $1", threadID).Scan()
}
<commit_msg>style: better sql formatting<commit_after>package main

import (
	//DB
	"database/sql"

	"github.com/golang/glog"
	"github.com/iambc/xerrors"
	_ "github.com/lib/pq"
)

//Relational database implementation for writer interface
type writerrdb struct {
	*sql.DB
}

func (db *writerrdb) getBoards(apiKey string) (currBoards []boards, err error) {
	glog.Info(" apiKey: ", apiKey)
	rows, err := db.Query("select b.ID, b.name, b.descr from boards b join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID where api_key = $1;", apiKey)
	if err != nil {
		return currBoards, xerrors.NewUIErr(err.Error(), err.Error(), `002`, true)
	}
	defer rows.Close()

	for rows.Next() {
		var board boards
		err = rows.Scan(&board.ID, &board.Name, &board.Descr)
		if err != nil {
			return currBoards, xerrors.NewUIErr(err.Error(), err.Error(), `003`, true)
		}
		currBoards = append(currBoards, board)
	}
	return currBoards, nil
}

func (db *writerrdb) getActiveThreadsForBoard(apiKey string, boardID int) (activeThreads []threads, err error) {
	rows, err := db.Query(`select t.ID, t.name, count(*), (select count(*) from thread_posts where thread_ID = t.ID and attachment_url is not null) from threads t
				join boards b on b.ID = t.board_ID
				join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID
				left join thread_posts tp on tp.thread_ID = t.ID
			    where t.is_active = TRUE and t.board_ID = $1 and ibc.api_key = $2 group by 1,2 order by t.ID;`, boardID, apiKey)
	if err != nil {
		return activeThreads, xerrors.NewUIErr(err.Error(), err.Error(), `006`, true)
	}
	defer rows.Close()

	for rows.Next() {
		var thread threads
		err = rows.Scan(&thread.ID, &thread.Name, &thread.PostCount, &thread.PostCountWithAttachment)
		if err != nil {
			return activeThreads, xerrors.NewUIErr(err.Error(), err.Error(), `007`, true)
		}
		activeThreads = append(activeThreads, thread)
	}
	return activeThreads, nil
}

func (db *writerrdb) getPostsForThread(apiKey string, threadID int) (currPosts []threadPosts, err error) {
	rows, err := db.Query(`select tp.ID, tp.body, tp.attachment_url, tp.inserted_at, tp.source_ip
			    from thread_posts tp join threads t on t.ID = tp.thread_ID
						 join boards b on b.ID = t.board_ID
						 join image_board_clusters ibc on ibc.ID = b.image_board_cluster_ID
			    where tp.thread_ID = $1 and ibc.api_key = $2 and t.is_active = true;`, threadID, apiKey)
	if err != nil {
		glog.Error(err)
		return currPosts, xerrors.NewSysErr()
	}
	defer rows.Close()

	for rows.Next() {
		var currPost threadPosts
		err = rows.Scan(&currPost.ID, &currPost.Body, &currPost.attachmentURL, &currPost.InsertedAt, &currPost.SourceIP)
		if err != nil {
			glog.Error(err)
			return currPosts, xerrors.NewSysErr()
		}
		currPosts = append(currPosts, currPost)
	}
	return currPosts, err
}

func (db *writerrdb) addPostToThread(threadID int, threadBodyPost string, attachmentURL *string, clientRemoteAddr string) (err error) {
	_, err = db.Query("INSERT INTO thread_posts(body, thread_ID, attachment_url, source_ip) VALUES($1, $2, $3, $4)", threadBodyPost, threadID, attachmentURL, clientRemoteAddr)

	if err != nil {
		glog.Error(err)
		return xerrors.NewUIErr(err.Error(), err.Error(), `011`, true)
	}
	return nil
}

func (db *writerrdb) addThread(boardID int, threadName string) (threads, error) {

	var threadID int
	err := db.QueryRow("INSERT INTO threads(name, board_ID, limits_reached_action_ID, max_posts_per_thread) VALUES($1, $2, 1, 10)  RETURNING ID, name", threadName, boardID).Scan(&threadID, &threadName)

	if err != nil {
		glog.Error("INSERT FAILED")
		return threads{ID: -1, Name: `err`}, xerrors.NewUIErr(err.Error(), err.Error(), `017`, true)
	}
	return threads{ID: threadID, Name: threadName}, nil
}

func (db *writerrdb) isThreadLimitReached(boardID int) (bool, error) {
	var isLimitReached bool
	err := db.QueryRow("select (select count(*) from threads  where board_ID = $1) > thread_setting_max_thread_count  from boards where ID = $1;", boardID).Scan(&isLimitReached)
	if err != nil {
		glog.Error("COULD NOT SELECT thread_count")
		return true, xerrors.NewUIErr(err.Error(), err.Error(), `015`, true)
	}

	return isLimitReached, nil
}

func (db *writerrdb) isPostLimitReached(threadID int) (bool, threads, error) {
	var isLimitReached bool
	var thread threads
	err := db.QueryRow("select (select count(*) from thread_posts  where thread_id = $1) > max_posts_per_thread, min_post_length, max_post_length  from threads where id = $1;", threadID).Scan(&isLimitReached, &thread.MinPostLength, &thread.MaxPostLength)
	if err != nil {
		return true, thread, xerrors.NewUIErr(err.Error(), err.Error(), `009`, true)
	}
	return isLimitReached, thread, err
}

func (db *writerrdb) archiveThread(threadID int) {
	db.QueryRow("UPDATE threads set is_active = false where ID = $1", threadID).Scan()
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()
	Log.Infoc(func() string {
		return fmt.Sprintf("conn is ok ?= " + this.conn != nil)
	})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	b, ok := this.in.ReadBuffer()
	for ; !ok; b, ok = this.in.ReadBuffer() {
		runtime.Gosched()
	}
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(*b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}
	for ok := this.out.WriteBuffer(&b); !ok; ok = this.out.WriteBuffer(&b) {
		runtime.Gosched()
	}
	this.outStat.increment(int64(1))

	return nil
}
<commit_msg>增加测试日志<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
	"runtime"
	"strconv"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(receiver): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			if err != nil {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
				})
				//         if err != io.EOF {
				//         }
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic(sender): %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()
	Log.Infoc(func() string {
		return fmt.Sprintf("conn is ok ?= " + strconv.FormatBool(this.conn != nil))
	})
	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, error) {
	var (
		err error
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, err
	}

	b, ok := this.in.ReadBuffer()
	for ; !ok; b, ok = this.in.ReadBuffer() {
		runtime.Gosched()
	}
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s)NewMessage  Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	_, err = msg.Decode(*b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Decode Error processing: %v", this.cid(), err)
		})
		return nil, err
	}

	return msg, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (error) {
	if this.out == nil {
		return ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	b := make([]byte, msg.Len())
	_, err := msg.Encode(b)
	if err != nil {
		return err
	}
	for ok := this.out.WriteBuffer(&b); !ok; ok = this.out.WriteBuffer(&b) {
		runtime.Gosched()
	}
	this.outStat.increment(int64(1))

	return nil
}
<|endoftext|>


<commit_before>/**
  *  _               _   _
  * | |             | | (_)
  * | |__   __ _ ___| |_ _  ___
  * | '_ \ / _` / __| __| |/ _ \
  * | | | | (_| \__ \ |_| |  __/
  * |_| |_|\__,_|___/\__|_|\___|
  *
  * Hastie - Static Site Generator
  *
  */

package main

import (
  "bytes"
  "conf"   /* from http://code.google.com/p/goconf/ */
  /* "goconf.googlecode.com/hg"*/
  "fmt"
  "github.com/russross/blackfriday"
  "io/ioutil"
  "os"
  "path/filepath"
  "strings"
  "sort"
  "template"
  "time"
)

const debug = false

type Page struct {
  Content     string
  Title       string
  Category    string
  Layout      string
  Pages       PagesSlice
  Recent      PagesSlice
  Date        *time.Time
  OutFile     string
}

type Config struct {
  TemplateDir string
  LayoutDir   string
  PublishDir  string
}

type PagesSlice []Page
func (p PagesSlice) Len() int           { return len(p) }
func (p PagesSlice) Less(i, j int) bool { return p[i].Date.Seconds() < p[j].Date.Seconds() }
func (p PagesSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p PagesSlice) Sort()              { sort.Sort(p) }

func main() {
  //TODO: add verbose / debug - command-line arguments

  config := setupConfig()

  // create list of directories and files
  site := &SiteStruct{}
  filepath.Walk(config.TemplateDir, site, nil)

  // create directories
  ///TODO: create necessary directories

  /* ******************************************
   * Loop through directories and build pages
   * ****************************************** */
  var pages PagesSlice
  for _,dir := range site.Directories {

    fmt.Println("Reading Directory: ", dir)
    directory_file_count := 0
    readglob := fmt.Sprintf("%s/*.md", dir)
    var dirfiles, _ = filepath.Glob(readglob)

    // loop through files in directory
    for _, file := range dirfiles {
      if debug { fmt.Println("  File: ", file) }
      outfile := filepath.Base(file)
      outfile = strings.Replace(outfile, ".md", ".html", 1)

      // read & parse file for parameters
      page := readParseFile(file)

      // skip file if no content
      if page.Content == "" {
        continue // skip to next file
      }
      pages = append(pages, page)
      directory_file_count += 1
    }
  }

  if (debug) { // spit out pages structure
    fmt.Println("################################################################################")
    fmt.Printf(" %-50s | %-10s | %s \n", "Title", "Category", "Outfile")
    fmt.Println("--------------------------------------------------------------------------------")
    for _, page := range pages {
      fmt.Printf(" %-50s | %-10s | %s \n", page.Title, page.Category, page.OutFile)
    }
    fmt.Println("################################################################################")
  }

  /* ******************************************
   * Create any data needed from pages
   * for example recent file list
   * category list, etc...
   * ****************************************** */

  // build recent file list, sorted
  recentList := getRecentList(pages)

  // Next Loop through Pages and apply templates and write out
  for _, page := range pages {

    fmt.Println("  Generating Template: ", page.OutFile)

    /* Assign global data to page object
     * Note: need better templating duplicating data
             since no logic in templates to limit to 3 */
    page.Pages = recentList
    if len(recentList) > 3 {
      page.Recent = recentList[0:3]
    } else {
      page.Recent = recentList
    }

    /* apply template - writes page data to buffer */
    buffer := new(bytes.Buffer)
    // read and parse all templates, since header, footer need
    // to be read and parsed for them to be included
    layoutsglob := fmt.Sprintf("%s/*.html", config.LayoutDir)
    ts, err := template.ParseTemplateGlob(layoutsglob)
    if err != nil {
      fmt.Println("Error Parsing Templates: ", err)
      os.Exit(1)
    }
    // pick layout based on specified
    templateFile := ""
    if page.Layout == "" {
      templateFile = "post.html"
    } else {
      templateFile = fmt.Sprintf("%s.html", page.Layout)
    }
    ts.Execute(buffer, templateFile, page)

    // writing out file
    writedir := fmt.Sprintf("%s/%s", config.PublishDir, page.Category)
    if (debug) { fmt.Println(" Write Directory: ", writedir) }
    os.MkdirAll(writedir, 0755) // does nothing if already exists

    outfile := fmt.Sprintf("%s/%s", config.PublishDir, page.OutFile)
    if (debug) { fmt.Println(" Writing File: ", outfile) }
    ioutil.WriteFile(outfile, []byte(buffer.String()), 0644)
  }

}

/* ************************************************
 * Read and Parse File
 * @param filename
 * @return content, params
 *
 * Note: named return params allows return to "know" what to return
 * ************************************************ */
func readParseFile(filename string) (page Page) {
  if (debug) { fmt.Println("Parsing File: ", filename) }
  epoch, _ := time.Parse("20060102", "19700101")
  // setup default page struct
  page = Page{
    Title: "",
    Category: "",
    Content: "",
    Layout: "",
    Date: epoch,
    OutFile: filename }

  // read file
  var data, err = ioutil.ReadFile(filename)
  if err != nil {
    fmt.Println("Error Reading: ", filename)
    return
  }

  // go through content parse from --- to ---
  var lines = strings.Split(string(data), "\n")
  var found = 0
  for i, line := range lines {
    line = strings.TrimSpace(line)

    if found == 1 {
      // parse line for param
      colonIndex := strings.Index(line, ":")
      if colonIndex > 0 {
        key := strings.TrimSpace(line[:colonIndex])
        value := strings.TrimSpace(line[colonIndex+1:])
        switch key {
          case "title": page.Title = value
          case "category": page.Category = value
          case "layout": page.Layout = value
        }
      }

    } else if found >= 2 {
      // params over
      lines = lines[i:]
      break
    }

    if line == "---" {
      found += 1
    }

  }

  // switch directory name to just category

  // chop off first directory, since that is the template dir
  page.OutFile = filename[strings.Index(filename, "/")+1:]
  page.OutFile = strings.Replace(page.OutFile, ".md", ".html", 1)

  // first directory is category
  if strings.Contains(page.OutFile, "/") {
    page.Category = page.OutFile[0:strings.Index(page.OutFile, "/")]
  }

  // parse date from filename, first part is just sample
  base := filepath.Base(page.OutFile)
  if base[0:2] == "20" || base[0:2] == "19" {  //HACK: if file starts with 20 assume date
    page.Date, _ = time.Parse("2006-01-02", base[0:10])
    page.OutFile = strings.Replace(page.OutFile, base[0:11], "", 1)  // remove date from final filename
  }

  // convert markdown content
  content := strings.Join(lines, "\n")
  output := blackfriday.MarkdownCommon([]byte(content))
  page.Content = string(output)

  return page
}

/* ************************************************
 * Build Recent File List
 *    - return array sorted most recent first
 *    - array includes real link (no date)
 *    - does not include files without date
 * ************************************************ */
func getRecentList(pages PagesSlice) (list PagesSlice) {
  fmt.Println("Creating Recent File List")
  for _, page := range pages {
    if (page.Date.Format("2006") != "1970") {
      list = append(list, page)
    }
  }
  list.Sort()

  // reverse
  for i, j := 0, len(list)-1; i < j; i, j = i+1, j-1 {
    list[i], list[j] = list[j], list[i]
  }

  return list
}

/* ************************************************
 * Takes a string array and sorts then reverses
 * Used for creating a most recent list
 * ************************************************ */
func sortReverse(arr []string) []string {
  ss := sort.StringSlice(arr)
  ss.Sort()
  for i, j := 0, len(ss)-1; i < j; i, j = i+1, j-1 {
    ss[i], ss[j] = ss[j], ss[i]
  }
  return ss
}

/* ************************************************
 *  File / Directory Walker
 * ************************************************ */
type SiteStruct struct{
  Files []string
  Directories []string
  Categories []string
}

func (v *SiteStruct) VisitDir(path string, f *os.FileInfo) bool {
  v.Categories = append(v.Categories, f.Name)
  v.Directories = append(v.Directories, path)
  /* fmt.Println("Directory: ", f.Name)*/
  return true
}

func (v *SiteStruct) VisitFile(path string, f *os.FileInfo) {
  v.Files = append(v.Files, path)
  /* fmt.Printf("File: %s  -- Path: %s \n", f.Name, path)*/
}

/* ************************************************
 * Check if File / Directory Exists
 * ************************************************ */
func exists(path string) bool {
  _, err := os.Stat(path)
  if err != nil {
    if e, ok := err.(*os.PathError); ok && (e.Error == os.ENOENT || e.Error == os.ENOTDIR) {
      return false
    }
    return false
  }
  return true
}

/* ************************************************
 * Read Config file or set defaults
 * ************************************************ */
func setupConfig() (config Config){
  cfg, err := conf.ReadConfigFile("hastie.config")
  if err != nil {
    // set defaults
    config.TemplateDir = "posts"
    config.LayoutDir = "layouts"
    config.PublishDir = "public"
  } else {
    config.TemplateDir, _ = cfg.GetString("default", "template_dir")
    config.LayoutDir, _ = cfg.GetString("default", "layout_dir")
    config.PublishDir, _ = cfg.GetString("default", "publish_dir")
    config.TemplateDir = strings.TrimSpace(config.TemplateDir)
    config.LayoutDir = strings.TrimSpace(config.LayoutDir)
    config.PublishDir = strings.TrimSpace(config.PublishDir)
  }
  return config
}
<commit_msg>clean up comments in source<commit_after>/**
  *  _               _   _
  * | |             | | (_)
  * | |__   __ _ ___| |_ _  ___
  * | '_ \ / _` / __| __| |/ _ \
  * | | | | (_| \__ \ |_| |  __/
  * |_| |_|\__,_|___/\__|_|\___|
  *
  * Hastie - Static Site Generator
  *
  */

package main

import (
  "bytes"
  "conf"   /* from http://code.google.com/p/goconf/ */
  /* "goconf.googlecode.com/hg"*/
  "fmt"
  "github.com/russross/blackfriday"
  "io/ioutil"
  "os"
  "path/filepath"
  "strings"
  "sort"
  "template"
  "time"
)

const debug = false

type Page struct {
  Content     string
  Title       string
  Category    string
  Layout      string
  Pages       PagesSlice
  Recent      PagesSlice
  Date        *time.Time
  OutFile     string
}

type Config struct {
  TemplateDir string
  LayoutDir   string
  PublishDir  string
}

type PagesSlice []Page
func (p PagesSlice) Len() int           { return len(p) }
func (p PagesSlice) Less(i, j int) bool { return p[i].Date.Seconds() < p[j].Date.Seconds() }
func (p PagesSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p PagesSlice) Sort()              { sort.Sort(p) }

func main() {
  //TODO: add verbose / debug - command-line arguments

  config := setupConfig()

  // create list of directories and files
  site := &SiteStruct{}
  filepath.Walk(config.TemplateDir, site, nil)

  /* ******************************************
   * Loop through directories and build pages
   * ****************************************** */
  var pages PagesSlice
  for _,dir := range site.Directories {

    fmt.Println("Reading Directory: ", dir)
    directory_file_count := 0
    readglob := fmt.Sprintf("%s/*.md", dir)
    var dirfiles, _ = filepath.Glob(readglob)

    // loop through files in directory
    for _, file := range dirfiles {
      if debug { fmt.Println("  File: ", file) }
      outfile := filepath.Base(file)
      outfile = strings.Replace(outfile, ".md", ".html", 1)

      // read & parse file for parameters
      page := readParseFile(file)

      // skip file if no content
      if page.Content == "" {
        continue // skip to next file
      }
      pages = append(pages, page)
      directory_file_count += 1
    }
  }

  if (debug) { // spit out pages structure
    fmt.Println("################################################################################")
    fmt.Printf(" %-50s | %-10s | %s \n", "Title", "Category", "Outfile")
    fmt.Println("--------------------------------------------------------------------------------")
    for _, page := range pages {
      fmt.Printf(" %-50s | %-10s | %s \n", page.Title, page.Category, page.OutFile)
    }
    fmt.Println("################################################################################")
  }

  /* ******************************************
   * Create any data needed from pages
   * for example recent file list
   * category list, etc...
   * ****************************************** */

  // build recent file list, sorted
  recentList := getRecentList(pages)

  /* ******************************************
   * Loop through pages and generate templates
   * ****************************************** */
  for _, page := range pages {

    fmt.Println("  Generating Template: ", page.OutFile)

    /* Assign global data to page object
     * Note: need better templating duplicating data
             since no logic in templates to limit to 3 */
    page.Pages = recentList
    if len(recentList) > 3 {
      page.Recent = recentList[0:3]
    } else {
      page.Recent = recentList
    }

    /* Templating - writes page data to buffer
     * read and parse all template files          */
    buffer := new(bytes.Buffer)
    layoutsglob := fmt.Sprintf("%s/*.html", config.LayoutDir)
    ts, err := template.ParseTemplateGlob(layoutsglob)
    if err != nil {
      fmt.Println("Error Parsing Templates: ", err)
      os.Exit(1)
    }
    // pick layout based on specified in file
    templateFile := ""
    if page.Layout == "" {
      templateFile = "post.html"
    } else {
      templateFile = fmt.Sprintf("%s.html", page.Layout)
    }
    ts.Execute(buffer, templateFile, page)

    // writing out file
    writedir := fmt.Sprintf("%s/%s", config.PublishDir, page.Category)
    if (debug) { fmt.Println(" Write Directory: ", writedir) }
    os.MkdirAll(writedir, 0755) // does nothing if already exists

    outfile := fmt.Sprintf("%s/%s", config.PublishDir, page.OutFile)
    if (debug) { fmt.Println(" Writing File: ", outfile) }
    ioutil.WriteFile(outfile, []byte(buffer.String()), 0644)
  }

}

/* ************************************************
 * Read and Parse File
 * @param filename
 * @return Page object
 * ************************************************ */
func readParseFile(filename string) (page Page) {
  if (debug) { fmt.Println("Parsing File: ", filename) }
  epoch, _ := time.Parse("20060102", "19700101")

  // setup default page struct
  page = Page{
    Title: "",
    Category: "",
    Content: "",
    Layout: "",
    Date: epoch,
    OutFile: filename }

  // read file
  var data, err = ioutil.ReadFile(filename)
  if err != nil {
    fmt.Println("Error Reading: ", filename)
    return
  }

  // go through content parse from --- to ---
  var lines = strings.Split(string(data), "\n")
  var found = 0
  for i, line := range lines {
    line = strings.TrimSpace(line)

    if found == 1 {
      // parse line for param
      colonIndex := strings.Index(line, ":")
      if colonIndex > 0 {
        key := strings.TrimSpace(line[:colonIndex])
        value := strings.TrimSpace(line[colonIndex+1:])
        switch key {
          case "title": page.Title = value
          case "category": page.Category = value
          case "layout": page.Layout = value
        }
      }

    } else if found >= 2 {
      // params over
      lines = lines[i:]
      break
    }

    if line == "---" {
      found += 1
    }

  }

  // switch directory name to just category

  // chop off first directory, since that is the template dir
  page.OutFile = filename[strings.Index(filename, "/")+1:]
  page.OutFile = strings.Replace(page.OutFile, ".md", ".html", 1)

  // next first directory is category
  if strings.Contains(page.OutFile, "/") {
    page.Category = page.OutFile[0:strings.Index(page.OutFile, "/")]
  }

  // parse date from filename
  base := filepath.Base(page.OutFile)
  if base[0:2] == "20" || base[0:2] == "19" {  //HACK: if file starts with 20 or 19 assume date
    page.Date, _ = time.Parse("2006-01-02", base[0:10])
    page.OutFile = strings.Replace(page.OutFile, base[0:11], "", 1)  // remove date from final filename
  }

  // convert markdown content
  content := strings.Join(lines, "\n")
  output := blackfriday.MarkdownCommon([]byte(content))
  page.Content = string(output)

  return page
}

/* ************************************************
 * Build Recent File List
 *    - return array sorted most recent first
 *    - array includes real link (no date)
 *    - does not include files without date
 * ************************************************ */
func getRecentList(pages PagesSlice) (list PagesSlice) {
  fmt.Println("Creating Recent File List")
  for _, page := range pages {
    // pages without dates are set to epoch
    if (page.Date.Format("2006") != "1970") {
      list = append(list, page)
    }
  }
  list.Sort()

  // reverse
  for i, j := 0, len(list)-1; i < j; i, j = i+1, j-1 {
    list[i], list[j] = list[j], list[i]
  }

  return list
}

/* ************************************************
 *  File / Directory Walker
 * ************************************************ */
type SiteStruct struct{
  Files []string
  Directories []string
  Categories []string
}

func (v *SiteStruct) VisitDir(path string, f *os.FileInfo) bool {
  v.Categories = append(v.Categories, f.Name)
  v.Directories = append(v.Directories, path)
  /* fmt.Println("Directory: ", f.Name)*/
  return true
}

func (v *SiteStruct) VisitFile(path string, f *os.FileInfo) {
  v.Files = append(v.Files, path)
  /* fmt.Printf("File: %s  -- Path: %s \n", f.Name, path)*/
}

/* ************************************************
 * Check if File / Directory Exists
 * ************************************************ */
func exists(path string) bool {
  _, err := os.Stat(path)
  if err != nil {
    if e, ok := err.(*os.PathError); ok && (e.Error == os.ENOENT || e.Error == os.ENOTDIR) {
      return false
    }
    return false
  }
  return true
}

/* ************************************************
 * Read Config file or set defaults
 * ************************************************ */
func setupConfig() (config Config){
  cfg, err := conf.ReadConfigFile("hastie.config")
  if err != nil {
    // set defaults
    config.TemplateDir = "posts"
    config.LayoutDir = "layouts"
    config.PublishDir = "public"
  } else {
    config.TemplateDir, _ = cfg.GetString("default", "template_dir")
    config.LayoutDir, _ = cfg.GetString("default", "layout_dir")
    config.PublishDir, _ = cfg.GetString("default", "publish_dir")
    config.TemplateDir = strings.TrimSpace(config.TemplateDir)
    config.LayoutDir = strings.TrimSpace(config.LayoutDir)
    config.PublishDir = strings.TrimSpace(config.PublishDir)
  }
  return config
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/kballard/go-shellquote"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Debugf is a helper function for debug logging if global variable debug is set to true
func Debugf(s string) {
	if debug != false {
		log.Print("DEBUG " + fmt.Sprint(s))
	}
}

// Verbosef is a helper function for debug logging if global variable verbose is set to true
func Verbosef(s string) {
	if debug != false || verbose != false {
		log.Print(fmt.Sprint(s))
	}
}

// Infof is a helper function for debug logging if global variable info is set to true
func Infof(s string) {
	if debug != false || verbose != false || info != false {
		fmt.Println(s)
	}
}

// fileExists checks if the given file exists and return a bool
func fileExists(file string) bool {
	if _, err := os.Stat(file); os.IsNotExist(err) {
		return false
	} else {
		return true
	}
}

// checkDirAndCreate tests if the given directory exists and tries to create it
func checkDirAndCreate(dir string, name string) string {
	if len(dir) != 0 {
		if _, err := os.Stat(dir); os.IsNotExist(err) {
			//log.Printf("checkDirAndCreate(): trying to create dir '%s' as %s", dir, name)
			if err := os.MkdirAll(dir, 0777); err != nil {
				log.Print("checkDirAndCreate(): Error: failed to create directory: ", dir)
				os.Exit(1)
			}
		}
	} else {
		// TODO make dir optional
		log.Print("dir setting '" + name + "' missing! Exiting!")
		os.Exit(1)
	}
	if !strings.HasSuffix(dir, "/") {
		dir = dir + "/"
	}
	Debugf("Using as " + name + ": " + dir)
	return dir
}

func createOrPurgeDir(dir string, callingFunction string) {
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		Debugf("createOrPurgeDir(): Trying to create dir: " + dir + " called from " + callingFunction)
		os.Mkdir(dir, 0777)
	} else {
		Debugf("createOrPurgeDir(): Trying to remove: " + dir + " called from " + callingFunction)
		if err := os.RemoveAll(dir); err != nil {
			log.Print("createOrPurgeDir(): error: removing dir failed", err)
		}
		Debugf("createOrPurgeDir(): Trying to create dir: " + dir + " called from " + callingFunction)
		os.Mkdir(dir, 0777)
	}
}

func executeCommand(command string, timeout int, allowFail bool) ExecResult {
	Debugf("Executing " + command)
	parts := strings.SplitN(command, " ", 2)
	cmd := parts[0]
	cmdArgs := []string{}
	if len(parts) > 1 {
		args, err := shellquote.Split(parts[1])
		if err != nil {
			Debugf("executeCommand(): err: " + fmt.Sprint(err))
		} else {
			cmdArgs = args
		}
	}

	before := time.Now()
	out, err := exec.Command(cmd, cmdArgs...).CombinedOutput()
	duration := time.Since(before).Seconds()
	er := ExecResult{0, string(out)}
	if msg, ok := err.(*exec.ExitError); ok { // there is error code
		er.returnCode = msg.Sys().(syscall.WaitStatus).ExitStatus()
	}
	mutex.Lock()
	syncGitTime += duration
	mutex.Unlock()
	Verbosef("Executing " + command + " took " + strconv.FormatFloat(duration, 'f', 5, 64) + "s")
	if err != nil && !allowFail {
		log.Print("executeCommand(): git command failed: " + command + " " + fmt.Sprint(err))
		log.Print("executeCommand(): Output: " + string(out))
		os.Exit(1)
	}
	return er
}
<commit_msg>add missing GitLab deploy key hint<commit_after>package main

import (
	"fmt"
	"github.com/kballard/go-shellquote"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Debugf is a helper function for debug logging if global variable debug is set to true
func Debugf(s string) {
	if debug != false {
		log.Print("DEBUG " + fmt.Sprint(s))
	}
}

// Verbosef is a helper function for debug logging if global variable verbose is set to true
func Verbosef(s string) {
	if debug != false || verbose != false {
		log.Print(fmt.Sprint(s))
	}
}

// Infof is a helper function for debug logging if global variable info is set to true
func Infof(s string) {
	if debug != false || verbose != false || info != false {
		fmt.Println(s)
	}
}

// fileExists checks if the given file exists and return a bool
func fileExists(file string) bool {
	if _, err := os.Stat(file); os.IsNotExist(err) {
		return false
	} else {
		return true
	}
}

// checkDirAndCreate tests if the given directory exists and tries to create it
func checkDirAndCreate(dir string, name string) string {
	if len(dir) != 0 {
		if _, err := os.Stat(dir); os.IsNotExist(err) {
			//log.Printf("checkDirAndCreate(): trying to create dir '%s' as %s", dir, name)
			if err := os.MkdirAll(dir, 0777); err != nil {
				log.Print("checkDirAndCreate(): Error: failed to create directory: ", dir)
				os.Exit(1)
			}
		}
	} else {
		// TODO make dir optional
		log.Print("dir setting '" + name + "' missing! Exiting!")
		os.Exit(1)
	}
	if !strings.HasSuffix(dir, "/") {
		dir = dir + "/"
	}
	Debugf("Using as " + name + ": " + dir)
	return dir
}

func createOrPurgeDir(dir string, callingFunction string) {
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		Debugf("createOrPurgeDir(): Trying to create dir: " + dir + " called from " + callingFunction)
		os.Mkdir(dir, 0777)
	} else {
		Debugf("createOrPurgeDir(): Trying to remove: " + dir + " called from " + callingFunction)
		if err := os.RemoveAll(dir); err != nil {
			log.Print("createOrPurgeDir(): error: removing dir failed", err)
		}
		Debugf("createOrPurgeDir(): Trying to create dir: " + dir + " called from " + callingFunction)
		os.Mkdir(dir, 0777)
	}
}

func executeCommand(command string, timeout int, allowFail bool) ExecResult {
	Debugf("Executing " + command)
	parts := strings.SplitN(command, " ", 2)
	cmd := parts[0]
	cmdArgs := []string{}
	if len(parts) > 1 {
		args, err := shellquote.Split(parts[1])
		if err != nil {
			Debugf("executeCommand(): err: " + fmt.Sprint(err))
		} else {
			cmdArgs = args
		}
	}

	before := time.Now()
	out, err := exec.Command(cmd, cmdArgs...).CombinedOutput()
	duration := time.Since(before).Seconds()
	er := ExecResult{0, string(out)}
	if msg, ok := err.(*exec.ExitError); ok { // there is error code
		er.returnCode = msg.Sys().(syscall.WaitStatus).ExitStatus()
	}
	mutex.Lock()
	syncGitTime += duration
	mutex.Unlock()
	Verbosef("Executing " + command + " took " + strconv.FormatFloat(duration, 'f', 5, 64) + "s")
	if err != nil && !allowFail {
		log.Print("executeCommand(): git command failed: " + command + " " + fmt.Sprint(err))
		log.Print("executeCommand(): Output: " + string(out))
		log.Println("If you are using GitLab be sure that you added your deploy key to your repository")
		os.Exit(1)
	}
	return er
}
<|endoftext|>


<commit_before>package com

import (
	"github.com/jacobsantos/go-com/idispatch"
	"github.com/jacobsantos/go-com/iunknown"
)

// ClassIDFrom retrieves class ID whether given is program ID or application string.
func ClassIDFrom(programID string) (classID *GUID, err error) {
	classID, err = ClassIDFromProgramID(programID)
	if err != nil {
		classID, err = ClassIDFromString(programID)
		if err != nil {
			return
		}
	}
	return
}

// InterfaceIDFrom retrieves Interface ID from interface type.
//
// Only supports IUnknown and IDispatch.
//
// TODO: support more COM interfaces.
func InterfaceIDFrom(obj interface{}) *GUID {
	var interfaceID *GUID
	switch obj.(type) {
	case iunknown.Unknown, *iunknown.Unknown:
		interfaceID = IUnknownInterfaceID
	case idispatch.Dispatch, *idispatch.Dispatch:
		interfaceID = IDispatchInterfaceID
	default:
		interfaceID = IUnknownInterfaceID
	}

	return interfaceID
}

// LoadObject loads COM object if in list.
func LoadObject(obj interface{}, programID ...string) []error {
	return LoadObjectWithCallback(CreateObject, &obj, programID...)
}

// LoadActiveObject loads COM active object if in list.
func LoadActiveObject(obj interface{}, programID ...string) []error {
	return LoadObjectWithCallback(GetActiveObject, &obj, programID...)
}

// LoadObjectWithCallback loads COM object if in list with callback.
func LoadObjectWithCallback(callback func(programID string, obj interface{}) error, obj interface{}, programID ...string) (err []error) {
	tempErrors := make([]error, len(names))
	var numErrors int
	for _, name := range names {
		err := callback(name, &obj)
		if err != nil {
			tempErrors = append(tempErrors, err)
			numErrors += 1
			continue
		}
		break
	}

	copy(errors, tempErrors[0:numErrors])
	return
}

// CreateObject creates object from programID based on interface type.
//
// Only supports IUnknown and IDispatch.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func CreateObject(programID string, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = CoCreateInstance(classID, InterfaceIDFrom(&obj), &obj)
	if err != nil {
		return
	}
	return
}

// CreateObjectWithInterfaceID creates object from programID.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func CreateObjectWithInterfaceID(programID string, interfaceID *GUID, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = CoCreateInstance(classID, interfaceID, &obj)
	if err != nil {
		return
	}
	return
}

// GetActiveObject retrieves active object for program ID and interface ID based
// on interface type.
//
// Only supports IUnknown and IDispatch.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func GetActiveObject(programID string, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = GetActiveObject(classID, InterfaceIDFrom(&obj), &obj)
	if err != nil {
		return
	}
	return
}

// GetActiveObjectWithInterfaceID retrieves active object for program ID and
// interface ID.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func GetActiveObjectWithInterfaceID(programID string, interfaceID *GUID, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = GetActiveObject(classID, interfaceID, &obj)
	if err != nil {
		return
	}
	return
}
<commit_msg>Fix package reference.<commit_after>package com

import (
	"github.com/go-ole/types/idispatch"
	"github.com/go-ole/types/iunknown"
)

// ClassIDFrom retrieves class ID whether given is program ID or application string.
func ClassIDFrom(programID string) (classID *GUID, err error) {
	classID, err = ClassIDFromProgramID(programID)
	if err != nil {
		classID, err = ClassIDFromString(programID)
		if err != nil {
			return
		}
	}
	return
}

// InterfaceIDFrom retrieves Interface ID from interface type.
//
// Only supports IUnknown and IDispatch.
//
// TODO: support more COM interfaces.
func InterfaceIDFrom(obj interface{}) *GUID {
	var interfaceID *GUID
	switch obj.(type) {
	case iunknown.Unknown, *iunknown.Unknown:
		interfaceID = IUnknownInterfaceID
	case idispatch.Dispatch, *idispatch.Dispatch:
		interfaceID = IDispatchInterfaceID
	default:
		interfaceID = IUnknownInterfaceID
	}

	return interfaceID
}

// LoadObject loads COM object if in list.
func LoadObject(obj interface{}, programID ...string) []error {
	return LoadObjectWithCallback(CreateObject, &obj, programID...)
}

// LoadActiveObject loads COM active object if in list.
func LoadActiveObject(obj interface{}, programID ...string) []error {
	return LoadObjectWithCallback(GetActiveObject, &obj, programID...)
}

// LoadObjectWithCallback loads COM object if in list with callback.
func LoadObjectWithCallback(callback func(programID string, obj interface{}) error, obj interface{}, programID ...string) (err []error) {
	tempErrors := make([]error, len(names))
	var numErrors int
	for _, name := range names {
		err := callback(name, &obj)
		if err != nil {
			tempErrors = append(tempErrors, err)
			numErrors += 1
			continue
		}
		break
	}

	copy(errors, tempErrors[0:numErrors])
	return
}

// CreateObject creates object from programID based on interface type.
//
// Only supports IUnknown and IDispatch.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func CreateObject(programID string, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = CoCreateInstance(classID, InterfaceIDFrom(&obj), &obj)
	if err != nil {
		return
	}
	return
}

// CreateObjectWithInterfaceID creates object from programID.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func CreateObjectWithInterfaceID(programID string, interfaceID *GUID, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = CoCreateInstance(classID, interfaceID, &obj)
	if err != nil {
		return
	}
	return
}

// GetActiveObject retrieves active object for program ID and interface ID based
// on interface type.
//
// Only supports IUnknown and IDispatch.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func GetActiveObject(programID string, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = GetActiveObject(classID, InterfaceIDFrom(&obj), &obj)
	if err != nil {
		return
	}
	return
}

// GetActiveObjectWithInterfaceID retrieves active object for program ID and
// interface ID.
//
// Program ID can be either program ID or application string.
//
// Object needs to be passed by reference.
func GetActiveObjectWithInterfaceID(programID string, interfaceID *GUID, obj interface{}) (err error) {
	classID, err := ClassIDFrom(programID)
	if err != nil {
		return
	}

	err = GetActiveObject(classID, interfaceID, &obj)
	if err != nil {
		return
	}
	return
}
<|endoftext|>


<commit_before>package grains

import (
	"errors"
)

const testVersion = 1

func Square(n int) (uint64, error) {
	if n<1 || n > 64 {
		return 0, errors.New("Invalid square number")
	}

	return uint64(1) << uint(n-1), nil
}

func Total() uint64 {
	return ^uint64(0)
}<commit_msg>'grains' - applied `gofmt`<commit_after>package grains

import (
	"errors"
)

const testVersion = 1

func Square(n int) (uint64, error) {
	if n < 1 || n > 64 {
		return 0, errors.New("Invalid square number")
	}

	return uint64(1) << uint(n-1), nil
}

func Total() uint64 {
	return ^uint64(0)
}
<|endoftext|>


<commit_before>/* Copyright 2017 LinkedIn Corp. Licensed under the Apache License, Version
 * 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

package core

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/viper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gopkg.in/natefinch/lumberjack.v2"
)

// CheckAndCreatePidFile takes a single argument, which is the path to a PID file (a file that contains a single
// integer, which is the process ID of a running process). If this file exists, and if the PID is that of a running
// process, return false as that indicates another copy of this process is already running. Otherwise, create the
// file and write this process's PID to the file and return true. Any error doing this (such as not having permissions
// to write the file) will return false.
//
// This func should be called when Burrow starts to prevent multiple copies from running.
func CheckAndCreatePidFile(filename string) bool {
	// Check if the PID file exists
	if _, err := os.Stat(filename); !os.IsNotExist(err) {
		// The file exists, so read it and check if the PID specified is running
		pidString, err := ioutil.ReadFile(filename)
		if err != nil {
			fmt.Printf("Cannot read PID file: %v", err)
			return false
		}
		pid, err := strconv.Atoi(string(pidString))
		if err != nil {
			fmt.Printf("Cannot interpret contents of PID file: %v", err)
			return false
		}

		if pid == os.Getpid() {
			// This could happen inside a docker
			// container, e.g. the pid of Burrow could be
			// equal to 1 each time the container is
			// restarted.
			return true
		}

		// Try sending a signal to the process to see if it is still running
		process, err := os.FindProcess(int(pid))
		if err == nil {
			err = process.Signal(syscall.Signal(0))
			if (err == nil) || (err == syscall.EPERM) {
				// The process exists, so we're going to assume it's an old Burrow and we shouldn't start
				fmt.Printf("Existing process running on PID %d. Exiting", pid)
				return false
			}
		}

	}

	// Create a PID file, replacing any existing one (as we already checked it)
	pidfile, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Printf("Cannot write PID file: %v", err)
		return false
	}
	fmt.Fprintf(pidfile, "%v", os.Getpid())
	pidfile.Close()
	return true
}

// RemovePidFile takes a single argument, which is the path to a PID file. That file is deleted. This func should be
// called when Burrow exits.
func RemovePidFile(filename string) {
	err := os.Remove(filename)
	if err != nil {
		fmt.Printf("Failed to remove PID file: %v\n", err)
	}
}

// ConfigureLogger returns a configured zap.Logger which can be used by Burrow for all logging. It also returns a
// zap.AtomicLevel, which can be used to dynamically adjust the level of the logger. The configuration for the logger
// is read from viper, with the following defaults:
//
// logging.level = info
//
// If logging.filename (path to the log file) is provided, a rolling log file is set up using lumberjack. The
// configuration for that log file is read from viper, with the following defaults:
//
// logging.maxsize = 100
// logging.maxbackups = 10
// logging.maxage = 30
// logging.use-localtime = false
// logging.use-compression = false
func ConfigureLogger() (*zap.Logger, *zap.AtomicLevel) {
	var level zap.AtomicLevel
	var syncOutput zapcore.WriteSyncer

	// Set config defaults for logging
	viper.SetDefault("logging.level", "info")
	viper.SetDefault("logging.maxsize", 100)
	viper.SetDefault("logging.maxbackups", 10)
	viper.SetDefault("logging.maxage", 30)

	// Create an AtomicLevel that we can use elsewhere to dynamically change the logging level
	logLevel := viper.GetString("logging.level")
	switch strings.ToLower(logLevel) {
	case "", "info":
		level = zap.NewAtomicLevelAt(zap.InfoLevel)
	case "debug":
		level = zap.NewAtomicLevelAt(zap.DebugLevel)
	case "warn":
		level = zap.NewAtomicLevelAt(zap.WarnLevel)
	case "error":
		level = zap.NewAtomicLevelAt(zap.ErrorLevel)
	case "panic":
		level = zap.NewAtomicLevelAt(zap.PanicLevel)
	case "fatal":
		level = zap.NewAtomicLevelAt(zap.FatalLevel)
	default:
		fmt.Printf("Invalid log level supplied. Defaulting to info: %s", logLevel)
		level = zap.NewAtomicLevelAt(zap.InfoLevel)
	}

	// If a filename has been set, set up a rotating logger. Otherwise, use Stdout
	logFilename := viper.GetString("logging.filename")
	if logFilename != "" {
		syncOutput = zapcore.AddSync(&lumberjack.Logger{
			Filename:   logFilename,
			MaxSize:    viper.GetInt("logging.maxsize"),
			MaxBackups: viper.GetInt("logging.maxbackups"),
			MaxAge:     viper.GetInt("logging.maxage"),
			LocalTime:  viper.GetBool("logging.use-localtime"),
			Compress:   viper.GetBool("logging.use-compression"),
		})
	} else {
		syncOutput = zapcore.Lock(os.Stdout)
	}

	core := zapcore.NewCore(
		zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
		syncOutput,
		level,
	)
	logger := zap.New(core)
	zap.ReplaceGlobals(logger)
	return logger, &level
}

// OpenOutLog takes a single argument, which is the path to a log file. This process's stdout and stderr are redirected
// to this log file. The os.File object is returned so that it can be managed.
func OpenOutLog(filename string) *os.File {
	// Move existing out file to a dated file if it exists
	if _, err := os.Stat(filename); err == nil {
		if err = os.Rename(filename, filename+"."+time.Now().Format("2006-01-02_15:04:05")); err != nil {
			fmt.Printf("Cannot move old out file: %v", err)
			os.Exit(1)
		}
	}

	// Redirect stdout and stderr to out file
	logFile, _ := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_SYNC, 0644)
	internalDup2(logFile.Fd(), 1)
	internalDup2(logFile.Fd(), 2)
	return logFile
}
<commit_msg>debug pid file detection logic<commit_after>/* Copyright 2017 LinkedIn Corp. Licensed under the Apache License, Version
 * 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

package core

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/viper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gopkg.in/natefinch/lumberjack.v2"
)

// CheckAndCreatePidFile takes a single argument, which is the path to a PID file (a file that contains a single
// integer, which is the process ID of a running process). If this file exists, and if the PID is that of a running
// process, return false as that indicates another copy of this process is already running. Otherwise, create the
// file and write this process's PID to the file and return true. Any error doing this (such as not having permissions
// to write the file) will return false.
//
// This func should be called when Burrow starts to prevent multiple copies from running.
func CheckAndCreatePidFile(filename string) bool {
	// Check if the PID file exists
	if _, err := os.Stat(filename); !os.IsNotExist(err) {
		// The file exists, so read it and check if the PID specified is running
		pidString, err := ioutil.ReadFile(filename)
		if err != nil {
			fmt.Printf("Cannot read PID file: %v", err)
			return false
		}
		pid, err := strconv.Atoi(string(pidString))
		if err != nil {
			fmt.Printf("Cannot interpret contents of PID file: %v", err)
			return false
		}

		if pid == os.Getpid() {
			// This could happen inside a docker
			// container, e.g. the pid of Burrow could be
			// equal to 1 each time the container is
			// restarted.
			fmt.Println("Found existing pidfile matching current pid")
			return true
		}

		// Try sending a signal to the process to see if it is still running
		process, err := os.FindProcess(int(pid))
		if err == nil {
			err = process.Signal(syscall.Signal(0))
			if (err == nil) || (err == syscall.EPERM) {
				// The process exists, so we're going to assume it's an old Burrow and we shouldn't start
				fmt.Printf("Existing process running on PID %d. Exiting (my pid = %d)", pid, os.Getpid())
				return false
			}
		}

	}

	// Create a PID file, replacing any existing one (as we already checked it)
	pidfile, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Printf("Cannot write PID file: %v", err)
		return false
	}
	fmt.Fprintf(pidfile, "%v", os.Getpid())
	pidfile.Close()
	return true
}

// RemovePidFile takes a single argument, which is the path to a PID file. That file is deleted. This func should be
// called when Burrow exits.
func RemovePidFile(filename string) {
	err := os.Remove(filename)
	if err != nil {
		fmt.Printf("Failed to remove PID file: %v\n", err)
	}
}

// ConfigureLogger returns a configured zap.Logger which can be used by Burrow for all logging. It also returns a
// zap.AtomicLevel, which can be used to dynamically adjust the level of the logger. The configuration for the logger
// is read from viper, with the following defaults:
//
// logging.level = info
//
// If logging.filename (path to the log file) is provided, a rolling log file is set up using lumberjack. The
// configuration for that log file is read from viper, with the following defaults:
//
// logging.maxsize = 100
// logging.maxbackups = 10
// logging.maxage = 30
// logging.use-localtime = false
// logging.use-compression = false
func ConfigureLogger() (*zap.Logger, *zap.AtomicLevel) {
	var level zap.AtomicLevel
	var syncOutput zapcore.WriteSyncer

	// Set config defaults for logging
	viper.SetDefault("logging.level", "info")
	viper.SetDefault("logging.maxsize", 100)
	viper.SetDefault("logging.maxbackups", 10)
	viper.SetDefault("logging.maxage", 30)

	// Create an AtomicLevel that we can use elsewhere to dynamically change the logging level
	logLevel := viper.GetString("logging.level")
	switch strings.ToLower(logLevel) {
	case "", "info":
		level = zap.NewAtomicLevelAt(zap.InfoLevel)
	case "debug":
		level = zap.NewAtomicLevelAt(zap.DebugLevel)
	case "warn":
		level = zap.NewAtomicLevelAt(zap.WarnLevel)
	case "error":
		level = zap.NewAtomicLevelAt(zap.ErrorLevel)
	case "panic":
		level = zap.NewAtomicLevelAt(zap.PanicLevel)
	case "fatal":
		level = zap.NewAtomicLevelAt(zap.FatalLevel)
	default:
		fmt.Printf("Invalid log level supplied. Defaulting to info: %s", logLevel)
		level = zap.NewAtomicLevelAt(zap.InfoLevel)
	}

	// If a filename has been set, set up a rotating logger. Otherwise, use Stdout
	logFilename := viper.GetString("logging.filename")
	if logFilename != "" {
		syncOutput = zapcore.AddSync(&lumberjack.Logger{
			Filename:   logFilename,
			MaxSize:    viper.GetInt("logging.maxsize"),
			MaxBackups: viper.GetInt("logging.maxbackups"),
			MaxAge:     viper.GetInt("logging.maxage"),
			LocalTime:  viper.GetBool("logging.use-localtime"),
			Compress:   viper.GetBool("logging.use-compression"),
		})
	} else {
		syncOutput = zapcore.Lock(os.Stdout)
	}

	core := zapcore.NewCore(
		zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()),
		syncOutput,
		level,
	)
	logger := zap.New(core)
	zap.ReplaceGlobals(logger)
	return logger, &level
}

// OpenOutLog takes a single argument, which is the path to a log file. This process's stdout and stderr are redirected
// to this log file. The os.File object is returned so that it can be managed.
func OpenOutLog(filename string) *os.File {
	// Move existing out file to a dated file if it exists
	if _, err := os.Stat(filename); err == nil {
		if err = os.Rename(filename, filename+"."+time.Now().Format("2006-01-02_15:04:05")); err != nil {
			fmt.Printf("Cannot move old out file: %v", err)
			os.Exit(1)
		}
	}

	// Redirect stdout and stderr to out file
	logFile, _ := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_SYNC, 0644)
	internalDup2(logFile.Fd(), 1)
	internalDup2(logFile.Fd(), 2)
	return logFile
}
<|endoftext|>


<commit_before>package services

import (
	"crypto/tls"
	"encoding/gob"
	"fmt"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/googollee/go-socket.io"
	"github.com/play-with-docker/play-with-docker/config"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/twinj/uuid"
)

var (
	sessionsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "sessions",
		Help: "Sessions",
	})
	clientsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "clients",
		Help: "Clients",
	})
	instancesGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "instances",
		Help: "Instances",
	})
)

func init() {
	prometheus.MustRegister(sessionsGauge)
	prometheus.MustRegister(clientsGauge)
	prometheus.MustRegister(instancesGauge)
}

var wsServer *socketio.Server

type Session struct {
	rw           sync.Mutex
	Id           string               `json:"id"`
	Instances    map[string]*Instance `json:"instances"`
	clients      []*Client            `json:"-"`
	CreatedAt    time.Time            `json:"created_at"`
	ExpiresAt    time.Time            `json:"expires_at"`
	scheduled    bool                 `json:"-"`
	ticker       *time.Ticker         `json:"-"`
	PwdIpAddress string               `json:"pwd_ip_address"`
	Ready        bool                 `json:"ready"`
	Stack        string               `json:"stack"`
}

type sessionBuilderWriter struct {
	session *Session
}

func (s *sessionBuilderWriter) Write(p []byte) (n int, err error) {
	wsServer.BroadcastTo(s.session.Id, "session builder out", string(p))
	return len(p), nil
}

func (s *Session) Lock() {
	s.rw.Lock()
}

func (s *Session) Unlock() {
	s.rw.Unlock()
}

func (s *Session) GetSmallestViewPort() ViewPort {
	minRows := s.clients[0].ViewPort.Rows
	minCols := s.clients[0].ViewPort.Cols

	for _, c := range s.clients {
		minRows = uint(math.Min(float64(minRows), float64(c.ViewPort.Rows)))
		minCols = uint(math.Min(float64(minCols), float64(c.ViewPort.Cols)))
	}

	return ViewPort{Rows: minRows, Cols: minCols}
}

func (s *Session) DeployStack() error {
	s.Lock()
	defer s.Unlock()

	if s.Ready {
		// a stack was already deployed on this session, just ignore
		return nil
	}

	s.setReady(false)
	i, err := NewInstance(s, InstanceConfig{})
	if err != nil {
		log.Printf("Error creating instance for stack [%s]: %s\n", s.Stack, err)
		return err
	}
	err = i.UploadFromURL("https://raw.githubusercontent.com/play-with-docker/stacks/master" + s.Stack)
	if err != nil {
		log.Printf("Error uploading stack file [%s]: %s\n", s.Stack, err)
		return err
	}

	_, err = Exec(i.Name, []string{"docker", "swarm", "init", "--advertise-addr", "eth0"})
	if err != nil {
		log.Printf("Error executing stack [%s]: %s\n", s.Stack, err)
		return err
	}
	w := sessionBuilderWriter{session: s}
	fileName := path.Base(s.Stack)
	code, err := ExecAttach(i.Name, []string{"docker", "stack", "deploy", "-c", "/var/run/pwd/uploads/" + fileName, "pwd"}, &w)
	if err != nil {
		log.Printf("Error executing stack [%s]: %s\n", s.Stack, err)
		return err
	}

	log.Printf("Stack execution finished with code %d\n", code)
	s.setReady(true)
	if err := saveSessionsToDisk(); err != nil {
		return err
	}
	return nil
}

func (s *Session) setReady(ready bool) {
	s.Ready = ready
	wsServer.BroadcastTo(s.Id, "session ready", s.Ready)
}

func (s *Session) AddNewClient(c *Client) {
	s.clients = append(s.clients, c)
	setGauges()
}

func (s *Session) SchedulePeriodicTasks() {
	if s.scheduled {
		return
	}

	go func() {
		s.scheduled = true

		s.ticker = time.NewTicker(1 * time.Second)
		for range s.ticker.C {
			var wg = sync.WaitGroup{}
			wg.Add(len(s.Instances))
			for _, ins := range s.Instances {
				var i *Instance = ins
				if i.dockerClient == nil {
					// Need to create client to the DinD docker daemon

					// We check if the client needs to use TLS
					var tlsConfig *tls.Config
					if len(i.Cert) > 0 && len(i.Key) > 0 {
						tlsConfig = tlsconfig.ClientDefault()
						tlsConfig.InsecureSkipVerify = true
						tlsCert, err := tls.X509KeyPair(i.Cert, i.Key)
						if err != nil {
							log.Println("Could not load X509 key pair: %v. Make sure the key is not encrypted", err)
							continue
						}
						tlsConfig.Certificates = []tls.Certificate{tlsCert}
					}

					transport := &http.Transport{
						DialContext: (&net.Dialer{
							Timeout:   1 * time.Second,
							KeepAlive: 30 * time.Second,
						}).DialContext}
					if tlsConfig != nil {
						transport.TLSClientConfig = tlsConfig
					}
					cli := &http.Client{
						Transport: transport,
					}
					c, err := client.NewClient(fmt.Sprintf("http://%s:2375", i.IP), api.DefaultVersion, cli, nil)
					if err != nil {
						log.Println("Could not connect to DinD docker daemon", err)
					} else {
						i.dockerClient = c
					}
				}
				go func() {
					defer wg.Done()
					for _, t := range periodicTasks {
						err := t.Run(i)
						if err != nil {
							if strings.Contains(err.Error(), "No such container") {
								log.Printf("Container for instance [%s] doesn't exist any more. Deleting from session.\n", i.IP)
								DeleteInstance(i.session, i)
							} else {
								log.Println(err)
							}
							break
						}
					}
				}()
			}
			wg.Wait()
			// broadcast all information
			for _, ins := range s.Instances {
				ins.Ports = UInt16Slice(ins.tempPorts)
				sort.Sort(ins.Ports)
				ins.tempPorts = []uint16{}

				wsServer.BroadcastTo(ins.session.Id, "instance stats", ins.Name, ins.Mem, ins.Cpu, ins.IsManager, ins.Ports)
			}
		}
	}()
}

var sessions map[string]*Session

func init() {
	sessions = make(map[string]*Session)
}

func CreateWSServer() *socketio.Server {
	server, err := socketio.NewServer(nil)
	if err != nil {
		log.Fatal(err)
	}
	wsServer = server
	return server
}

func CloseSessionAfter(s *Session, d time.Duration) {
	time.AfterFunc(d, func() {
		CloseSession(s)
	})
}

func CloseSession(s *Session) error {
	s.rw.Lock()
	defer s.rw.Unlock()

	if s.ticker != nil {
		s.ticker.Stop()
	}
	wsServer.BroadcastTo(s.Id, "session end")
	for _, c := range s.clients {
		c.so.Emit("disconnect")
	}
	log.Printf("Starting clean up of session [%s]\n", s.Id)
	for _, i := range s.Instances {
		err := DeleteInstance(s, i)
		if err != nil {
			log.Println(err)
			return err
		}
	}
	// Disconnect PWD daemon from the network
	if err := DisconnectNetwork("pwd", s.Id); err != nil {
		if !strings.Contains(err.Error(), "is not connected to the network") {
			log.Println("ERROR NETWORKING")
			return err
		}
	}
	log.Printf("Disconnected pwd from network [%s]\n", s.Id)
	if err := DeleteNetwork(s.Id); err != nil {
		if !strings.Contains(err.Error(), "not found") {
			log.Println(err)
			return err
		}
	}
	delete(sessions, s.Id)

	// We store sessions as soon as we delete one
	if err := saveSessionsToDisk(); err != nil {
		return err
	}
	setGauges()
	log.Printf("Cleaned up session [%s]\n", s.Id)
	return nil
}

var defaultDuration = 4 * time.Hour

func GetDuration(reqDur string) time.Duration {
	if reqDur != "" {
		if dur, err := time.ParseDuration(reqDur); err == nil && dur <= defaultDuration {
			return dur
		}
		return defaultDuration
	}

	envDur := os.Getenv("EXPIRY")
	if dur, err := time.ParseDuration(envDur); err == nil {
		return dur
	}

	return defaultDuration
}

func NewSession(duration time.Duration, stack string) (*Session, error) {
	s := &Session{}
	s.Id = uuid.NewV4().String()
	s.Instances = map[string]*Instance{}
	s.CreatedAt = time.Now()
	s.ExpiresAt = s.CreatedAt.Add(duration)
	if stack == "" {
		s.Ready = true
	}
	s.Stack = stack
	log.Printf("NewSession id=[%s]\n", s.Id)

	// Schedule cleanup of the session
	CloseSessionAfter(s, duration)

	if err := CreateNetwork(s.Id); err != nil {
		log.Println("ERROR NETWORKING")
		return nil, err
	}
	log.Printf("Network [%s] created for session [%s]\n", s.Id, s.Id)

	// Connect PWD daemon to the new network
	ip, err := ConnectNetwork(config.PWDContainerName, s.Id, "")
	if err != nil {
		log.Println("ERROR NETWORKING")
		return nil, err
	}
	s.PwdIpAddress = ip
	log.Printf("Connected %s to network [%s]\n", config.PWDContainerName, s.Id)

	// Schedule peridic tasks execution
	s.SchedulePeriodicTasks()

	sessions[s.Id] = s

	// We store sessions as soon as we create one so we don't delete new sessions on an api restart
	if err := saveSessionsToDisk(); err != nil {
		return nil, err
	}

	setGauges()
	return s, nil
}

func GetSession(sessionId string) *Session {
	s := sessions[sessionId]
	if s != nil {
		for _, instance := range s.Instances {
			if !instance.IsConnected() {
				instance.SetSession(s)
				go instance.Attach()
			}
		}

	}
	return s
}

func setGauges() {
	var ins float64
	var cli float64

	for _, s := range sessions {
		ins += float64(len(s.Instances))
		cli += float64(len(s.clients))
	}

	clientsGauge.Set(cli)
	instancesGauge.Set(ins)
	sessionsGauge.Set(float64(len(sessions)))
}

func LoadSessionsFromDisk() error {
	file, err := os.Open(config.SessionsFile)
	if err == nil {
		decoder := gob.NewDecoder(file)
		err = decoder.Decode(&sessions)

		if err != nil {
			return err
		}

		// schedule session expiration
		for _, s := range sessions {
			timeLeft := s.ExpiresAt.Sub(time.Now())
			CloseSessionAfter(s, timeLeft)

			// start collecting stats for every instance
			for _, i := range s.Instances {
				// wire the session back to the instance
				i.session = s
			}

			// Connect PWD daemon to the new network
			if s.PwdIpAddress == "" {
				log.Fatal("Cannot load stored sessions as they don't have the pwd ip address stored with them")
			}
			if _, err := ConnectNetwork(config.PWDContainerName, s.Id, s.PwdIpAddress); err != nil {
				if strings.Contains(err.Error(), "Could not attach to network") {
					log.Printf("Network for session [%s] doesn't exist. Removing all instances and session.", s.Id)
					CloseSession(s)
				} else {
					log.Println("ERROR NETWORKING", err)
					return err
				}
			} else {
				log.Printf("Connected %s to network [%s]\n", config.PWDContainerName, s.Id)

				// Schedule peridic tasks execution
				s.SchedulePeriodicTasks()
			}
		}
	}
	file.Close()
	setGauges()
	return err
}

func saveSessionsToDisk() error {
	rw.Lock()
	defer rw.Unlock()
	file, err := os.Create(config.SessionsFile)
	if err == nil {
		encoder := gob.NewEncoder(file)
		err = encoder.Encode(&sessions)
	}
	file.Close()
	return err
}
<commit_msg>Do everything in a single command and pull all images before deploying the stack to the swarm.<commit_after>package services

import (
	"crypto/tls"
	"encoding/gob"
	"fmt"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/googollee/go-socket.io"
	"github.com/play-with-docker/play-with-docker/config"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/twinj/uuid"
)

var (
	sessionsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "sessions",
		Help: "Sessions",
	})
	clientsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "clients",
		Help: "Clients",
	})
	instancesGauge = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "instances",
		Help: "Instances",
	})
)

func init() {
	prometheus.MustRegister(sessionsGauge)
	prometheus.MustRegister(clientsGauge)
	prometheus.MustRegister(instancesGauge)
}

var wsServer *socketio.Server

type Session struct {
	rw           sync.Mutex
	Id           string               `json:"id"`
	Instances    map[string]*Instance `json:"instances"`
	clients      []*Client            `json:"-"`
	CreatedAt    time.Time            `json:"created_at"`
	ExpiresAt    time.Time            `json:"expires_at"`
	scheduled    bool                 `json:"-"`
	ticker       *time.Ticker         `json:"-"`
	PwdIpAddress string               `json:"pwd_ip_address"`
	Ready        bool                 `json:"ready"`
	Stack        string               `json:"stack"`
}

type sessionBuilderWriter struct {
	session *Session
}

func (s *sessionBuilderWriter) Write(p []byte) (n int, err error) {
	wsServer.BroadcastTo(s.session.Id, "session builder out", string(p))
	return len(p), nil
}

func (s *Session) Lock() {
	s.rw.Lock()
}

func (s *Session) Unlock() {
	s.rw.Unlock()
}

func (s *Session) GetSmallestViewPort() ViewPort {
	minRows := s.clients[0].ViewPort.Rows
	minCols := s.clients[0].ViewPort.Cols

	for _, c := range s.clients {
		minRows = uint(math.Min(float64(minRows), float64(c.ViewPort.Rows)))
		minCols = uint(math.Min(float64(minCols), float64(c.ViewPort.Cols)))
	}

	return ViewPort{Rows: minRows, Cols: minCols}
}

func (s *Session) DeployStack() error {
	s.Lock()
	defer s.Unlock()

	if s.Ready {
		// a stack was already deployed on this session, just ignore
		return nil
	}

	s.setReady(false)
	i, err := NewInstance(s, InstanceConfig{})
	if err != nil {
		log.Printf("Error creating instance for stack [%s]: %s\n", s.Stack, err)
		return err
	}
	err = i.UploadFromURL("https://raw.githubusercontent.com/play-with-docker/stacks/master" + s.Stack)
	if err != nil {
		log.Printf("Error uploading stack file [%s]: %s\n", s.Stack, err)
		return err
	}

	fileName := path.Base(s.Stack)
	file := fmt.Sprintf("/var/run/pwd/uploads/%s", fileName)
	cmd := fmt.Sprintf("docker swarm init --advertise-addr eth0 && docker-compose -f %s pull && docker stack deploy -c %s pwd", file, file)

	w := sessionBuilderWriter{session: s}
	code, err := ExecAttach(i.Name, []string{"sh", "-c", cmd}, &w)
	if err != nil {
		log.Printf("Error executing stack [%s]: %s\n", s.Stack, err)
		return err
	}

	log.Printf("Stack execution finished with code %d\n", code)
	s.setReady(true)
	if err := saveSessionsToDisk(); err != nil {
		return err
	}
	return nil
}

func (s *Session) setReady(ready bool) {
	s.Ready = ready
	wsServer.BroadcastTo(s.Id, "session ready", s.Ready)
}

func (s *Session) AddNewClient(c *Client) {
	s.clients = append(s.clients, c)
	setGauges()
}

func (s *Session) SchedulePeriodicTasks() {
	if s.scheduled {
		return
	}

	go func() {
		s.scheduled = true

		s.ticker = time.NewTicker(1 * time.Second)
		for range s.ticker.C {
			var wg = sync.WaitGroup{}
			wg.Add(len(s.Instances))
			for _, ins := range s.Instances {
				var i *Instance = ins
				if i.dockerClient == nil {
					// Need to create client to the DinD docker daemon

					// We check if the client needs to use TLS
					var tlsConfig *tls.Config
					if len(i.Cert) > 0 && len(i.Key) > 0 {
						tlsConfig = tlsconfig.ClientDefault()
						tlsConfig.InsecureSkipVerify = true
						tlsCert, err := tls.X509KeyPair(i.Cert, i.Key)
						if err != nil {
							log.Println("Could not load X509 key pair: %v. Make sure the key is not encrypted", err)
							continue
						}
						tlsConfig.Certificates = []tls.Certificate{tlsCert}
					}

					transport := &http.Transport{
						DialContext: (&net.Dialer{
							Timeout:   1 * time.Second,
							KeepAlive: 30 * time.Second,
						}).DialContext}
					if tlsConfig != nil {
						transport.TLSClientConfig = tlsConfig
					}
					cli := &http.Client{
						Transport: transport,
					}
					c, err := client.NewClient(fmt.Sprintf("http://%s:2375", i.IP), api.DefaultVersion, cli, nil)
					if err != nil {
						log.Println("Could not connect to DinD docker daemon", err)
					} else {
						i.dockerClient = c
					}
				}
				go func() {
					defer wg.Done()
					for _, t := range periodicTasks {
						err := t.Run(i)
						if err != nil {
							if strings.Contains(err.Error(), "No such container") {
								log.Printf("Container for instance [%s] doesn't exist any more. Deleting from session.\n", i.IP)
								DeleteInstance(i.session, i)
							} else {
								log.Println(err)
							}
							break
						}
					}
				}()
			}
			wg.Wait()
			// broadcast all information
			for _, ins := range s.Instances {
				ins.Ports = UInt16Slice(ins.tempPorts)
				sort.Sort(ins.Ports)
				ins.tempPorts = []uint16{}

				wsServer.BroadcastTo(ins.session.Id, "instance stats", ins.Name, ins.Mem, ins.Cpu, ins.IsManager, ins.Ports)
			}
		}
	}()
}

var sessions map[string]*Session

func init() {
	sessions = make(map[string]*Session)
}

func CreateWSServer() *socketio.Server {
	server, err := socketio.NewServer(nil)
	if err != nil {
		log.Fatal(err)
	}
	wsServer = server
	return server
}

func CloseSessionAfter(s *Session, d time.Duration) {
	time.AfterFunc(d, func() {
		CloseSession(s)
	})
}

func CloseSession(s *Session) error {
	s.rw.Lock()
	defer s.rw.Unlock()

	if s.ticker != nil {
		s.ticker.Stop()
	}
	wsServer.BroadcastTo(s.Id, "session end")
	for _, c := range s.clients {
		c.so.Emit("disconnect")
	}
	log.Printf("Starting clean up of session [%s]\n", s.Id)
	for _, i := range s.Instances {
		err := DeleteInstance(s, i)
		if err != nil {
			log.Println(err)
			return err
		}
	}
	// Disconnect PWD daemon from the network
	if err := DisconnectNetwork("pwd", s.Id); err != nil {
		if !strings.Contains(err.Error(), "is not connected to the network") {
			log.Println("ERROR NETWORKING")
			return err
		}
	}
	log.Printf("Disconnected pwd from network [%s]\n", s.Id)
	if err := DeleteNetwork(s.Id); err != nil {
		if !strings.Contains(err.Error(), "not found") {
			log.Println(err)
			return err
		}
	}
	delete(sessions, s.Id)

	// We store sessions as soon as we delete one
	if err := saveSessionsToDisk(); err != nil {
		return err
	}
	setGauges()
	log.Printf("Cleaned up session [%s]\n", s.Id)
	return nil
}

var defaultDuration = 4 * time.Hour

func GetDuration(reqDur string) time.Duration {
	if reqDur != "" {
		if dur, err := time.ParseDuration(reqDur); err == nil && dur <= defaultDuration {
			return dur
		}
		return defaultDuration
	}

	envDur := os.Getenv("EXPIRY")
	if dur, err := time.ParseDuration(envDur); err == nil {
		return dur
	}

	return defaultDuration
}

func NewSession(duration time.Duration, stack string) (*Session, error) {
	s := &Session{}
	s.Id = uuid.NewV4().String()
	s.Instances = map[string]*Instance{}
	s.CreatedAt = time.Now()
	s.ExpiresAt = s.CreatedAt.Add(duration)
	if stack == "" {
		s.Ready = true
	}
	s.Stack = stack
	log.Printf("NewSession id=[%s]\n", s.Id)

	// Schedule cleanup of the session
	CloseSessionAfter(s, duration)

	if err := CreateNetwork(s.Id); err != nil {
		log.Println("ERROR NETWORKING")
		return nil, err
	}
	log.Printf("Network [%s] created for session [%s]\n", s.Id, s.Id)

	// Connect PWD daemon to the new network
	ip, err := ConnectNetwork(config.PWDContainerName, s.Id, "")
	if err != nil {
		log.Println("ERROR NETWORKING")
		return nil, err
	}
	s.PwdIpAddress = ip
	log.Printf("Connected %s to network [%s]\n", config.PWDContainerName, s.Id)

	// Schedule peridic tasks execution
	s.SchedulePeriodicTasks()

	sessions[s.Id] = s

	// We store sessions as soon as we create one so we don't delete new sessions on an api restart
	if err := saveSessionsToDisk(); err != nil {
		return nil, err
	}

	setGauges()
	return s, nil
}

func GetSession(sessionId string) *Session {
	s := sessions[sessionId]
	if s != nil {
		for _, instance := range s.Instances {
			if !instance.IsConnected() {
				instance.SetSession(s)
				go instance.Attach()
			}
		}

	}
	return s
}

func setGauges() {
	var ins float64
	var cli float64

	for _, s := range sessions {
		ins += float64(len(s.Instances))
		cli += float64(len(s.clients))
	}

	clientsGauge.Set(cli)
	instancesGauge.Set(ins)
	sessionsGauge.Set(float64(len(sessions)))
}

func LoadSessionsFromDisk() error {
	file, err := os.Open(config.SessionsFile)
	if err == nil {
		decoder := gob.NewDecoder(file)
		err = decoder.Decode(&sessions)

		if err != nil {
			return err
		}

		// schedule session expiration
		for _, s := range sessions {
			timeLeft := s.ExpiresAt.Sub(time.Now())
			CloseSessionAfter(s, timeLeft)

			// start collecting stats for every instance
			for _, i := range s.Instances {
				// wire the session back to the instance
				i.session = s
			}

			// Connect PWD daemon to the new network
			if s.PwdIpAddress == "" {
				log.Fatal("Cannot load stored sessions as they don't have the pwd ip address stored with them")
			}
			if _, err := ConnectNetwork(config.PWDContainerName, s.Id, s.PwdIpAddress); err != nil {
				if strings.Contains(err.Error(), "Could not attach to network") {
					log.Printf("Network for session [%s] doesn't exist. Removing all instances and session.", s.Id)
					CloseSession(s)
				} else {
					log.Println("ERROR NETWORKING", err)
					return err
				}
			} else {
				log.Printf("Connected %s to network [%s]\n", config.PWDContainerName, s.Id)

				// Schedule peridic tasks execution
				s.SchedulePeriodicTasks()
			}
		}
	}
	file.Close()
	setGauges()
	return err
}

func saveSessionsToDisk() error {
	rw.Lock()
	defer rw.Unlock()
	file, err := os.Create(config.SessionsFile)
	if err == nil {
		encoder := gob.NewEncoder(file)
		err = encoder.Encode(&sessions)
	}
	file.Close()
	return err
}
<|endoftext|>


<commit_before>// Package herots provide fast way to create TLS services: server and client.
//
// Explanation of the name: HERald Of The Swarm
//
// By the way - have a nice day :)
package herots

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
)

// Server - primary struct for server implementation.
type Server struct {
	options *Options
	certs   struct {
		Cert tls.Certificate
		pool struct {
			IsSet bool
			Pool  *x509.CertPool
		}
	}
	listener       net.Listener
	logDestination io.Writer
}

// predefined errors messages
const (
	LoadKeyPairError      = "herots: load key pair error"
	LoadClientCaCertError = "herots srv: load client CA cert error"
	StartServerError      = "herots srv: start tls server error"
	NoKeyPairLoad         = "herots: no load key pair (use LoadKeyPair func)"
	AcceptConnError       = "herots srv: connection accept error"
)

// Options - structure, which is used to configure a TLS server and client.
type Options struct {
	// Server host.
	//
	// Default: '127.0.0.1'.
	Host string

	// Server port.
	//
	// Default: '9000'.
	Port int

	// LogLevel provides the opportunity to choose the level of
	// information messages.
	// Each level includes the messages from the previous level.
	// 0 - no messages
	// 1 - notice
	// 2 - info
	// 3 - error
	//
	// Default: '0'.
	LogLevel int

	// LogDestination provides the opportunity to choose the own
	// destination for log messages (errors, info, etc).
	//
	// Default: 'os.Stdout'.
	LogDestination io.Writer

	// TLSAuthType - refer to http://golang.org/pkg/crypto/tls/#ClientAuthType
	// Default: tls.RequireAnyClientCert
	TLSAuthType tls.ClientAuthType
}

// NewServer - function for create Server struct
func NewServer(o *Options) *Server {
	s := &Server{}

	// check mandatory options
	if o.LogDestination == nil {
		o.LogDestination = os.Stdout
	}

	if o.Port == 0 {
		o.Port = 9000
	}

	if o.TLSAuthType == 0 {
		o.TLSAuthType = tls.RequireAnyClientCert
	}

	s.options = o

	return s
}

// log - function for print log messages
func (s *Server) log(m string, lvl int) {
	if s.options.LogLevel == 0 {
		return
	}

	if s.options.LogLevel <= lvl {
		fmt.Fprintf(s.logDestination, "herots srv: %s\n", m)
	}
}

// LoadKeyPair - function for load certificate and private key pair.
//
// Public/private key pair require as PEM encoded data.
func (s *Server) LoadKeyPair(cert, key []byte) error {
	// create cert pool
	s.certs.pool.Pool = x509.NewCertPool()

	// load keypair
	c, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.Cert = c

	// add cert to pool
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.pool.Pool.AddCert(ca)
	s.certs.pool.IsSet = true

	s.log("load key pair ok", 2)

	return nil
}

// AddClientCACert - function for adding client CA certificate to
// x509.CertPool (tls.Config.ClientCAs).
//
// By default server add cert from server public/private key pair (LoadKeyPair)
// to cert pool.
func (s *Server) AddClientCACert(cert []byte) error {
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadClientCaCertError, err)
	}
	s.certs.pool.Pool.AddCert(ca)

	s.log("load client CA cert ok", 2)

	return nil
}

// Accept - accept and return connections.
func (s *Server) Accept() (net.Conn, error) {
	conn, err := s.listener.Accept()
	if err != nil {
		s.log("accept conn error: "+err.Error(), 3)
		return conn, fmt.Errorf("%s: %v\n", AcceptConnError, err)
	}
	s.log("accepted conn from "+conn.RemoteAddr().String(), 2)
	return conn, nil
}

// Start - function for start server.
func (s *Server) Start() error {
	// load keypair check
	if len(s.certs.Cert.Certificate) == 0 {
		return fmt.Errorf("%s\n", NoKeyPairLoad)
	}

	config := tls.Config{
		ClientAuth:   s.options.TLSAuthType,
		Certificates: []tls.Certificate{s.certs.Cert},
		ClientCAs:    s.certs.pool.Pool,
		Rand:         rand.Reader,
	}

	service := s.options.Host + ":" + strconv.Itoa(s.options.Port)

	listener, err := tls.Listen("tcp", service, &config)
	if err != nil {
		return fmt.Errorf("%s: %v\n", StartServerError, err)
	}
	s.listener = listener

	s.log("listening on "+service, 2)

	return nil
}
<commit_msg>create external logger & refactor Server for new logger<commit_after>// Package herots provide fast way to create TLS services: server and client.
//
// Explanation of the name: HERald Of The Swarm
//
// By the way - have a nice day :)
package herots

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
)

////////////////////////////////////////////////////////////////////////////////
//                             Shared functions                               //
////////////////////////////////////////////////////////////////////////////////

// log - struct for internal log service
type log struct {
	LogLevel       int
	LogDestination io.Writer
}

func (l *log) Log(msg string, lvl int) {
	if l.LogLevel == 0 {
		return
	}

	if l.LogLevel <= lvl {
		fmt.Fprintf(l.LogDestination, "herots: %s\n", msg)
	}
}

////////////////////////////////////////////////////////////////////////////////
//                                  Server                                    //
////////////////////////////////////////////////////////////////////////////////

// Server - primary struct for server implementation.
type Server struct {
	options *Options
	certs   struct {
		Cert tls.Certificate
		pool struct {
			IsSet bool
			Pool  *x509.CertPool
		}
	}
	listener net.Listener
	logger   *log
}

// predefined errors messages
const (
	LoadKeyPairError      = "herots: load key pair error"
	LoadClientCaCertError = "herots srv: load client CA cert error"
	StartServerError      = "herots srv: start tls server error"
	NoKeyPairLoad         = "herots: no load key pair (use LoadKeyPair func)"
	AcceptConnError       = "herots srv: connection accept error"
)

// Options - structure, which is used to configure a TLS server and client.
type Options struct {
	// Server host.
	//
	// Default: '127.0.0.1'.
	Host string

	// Server port.
	//
	// Default: '9000'.
	Port int

	// LogLevel provides the opportunity to choose the level of
	// information messages.
	// Each level includes the messages from the previous level.
	// 0 - no messages
	// 1 - notice
	// 2 - info
	// 3 - error
	//
	// Default: '0'.
	LogLevel int

	// LogDestination provides the opportunity to choose the own
	// destination for log messages (errors, info, etc).
	//
	// Default: 'os.Stdout'.
	LogDestination io.Writer

	// TLSAuthType - refer to http://golang.org/pkg/crypto/tls/#ClientAuthType
	//
	// This option ignored for client implementation.
	//
	// Default: tls.RequireAnyClientCert
	TLSAuthType tls.ClientAuthType
}

// NewServer - function for create Server struct
func NewServer(o *Options) *Server {
	s := &Server{}

	// check mandatory options
	if o.LogDestination == nil {
		o.LogDestination = os.Stdout
	}

	if o.Port == 0 {
		o.Port = 9000
	}

	if o.TLSAuthType == 0 {
		o.TLSAuthType = tls.RequireAnyClientCert
	}

	l := &log{
		LogLevel:       o.LogLevel,
		LogDestination: o.LogDestination,
	}

	s.options = o
	s.logger = l

	return s
}

// LoadKeyPair - function for load certificate and private key pair.
//
// Public/private key pair require as PEM encoded data.
func (s *Server) LoadKeyPair(cert, key []byte) error {
	// create cert pool
	s.certs.pool.Pool = x509.NewCertPool()

	// load keypair
	c, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.Cert = c

	// add cert to pool
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.pool.Pool.AddCert(ca)
	s.certs.pool.IsSet = true

	s.logger.Log("load key pair ok", 2)

	return nil
}

// AddClientCACert - function for adding client CA certificate to
// x509.CertPool (tls.Config.ClientCAs).
//
// By default server add cert from server public/private key pair (LoadKeyPair)
// to cert pool.
func (s *Server) AddClientCACert(cert []byte) error {
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadClientCaCertError, err)
	}
	s.certs.pool.Pool.AddCert(ca)

	s.logger.Log("load client CA cert ok", 2)

	return nil
}

// Accept - accept and return connections.
func (s *Server) Accept() (net.Conn, error) {
	conn, err := s.listener.Accept()
	if err != nil {
		s.logger.Log("accept conn error: "+err.Error(), 3)
		return conn, fmt.Errorf("%s: %v\n", AcceptConnError, err)
	}
	s.logger.Log("accepted conn from "+conn.RemoteAddr().String(), 2)
	return conn, nil
}

// Start - function for start server.
func (s *Server) Start() error {
	// load keypair check
	if len(s.certs.Cert.Certificate) == 0 {
		return fmt.Errorf("%s\n", NoKeyPairLoad)
	}

	config := tls.Config{
		ClientAuth:   s.options.TLSAuthType,
		Certificates: []tls.Certificate{s.certs.Cert},
		ClientCAs:    s.certs.pool.Pool,
		Rand:         rand.Reader,
	}

	service := s.options.Host + ":" + strconv.Itoa(s.options.Port)

	listener, err := tls.Listen("tcp", service, &config)
	if err != nil {
		return fmt.Errorf("%s: %v\n", StartServerError, err)
	}
	s.listener = listener

	s.logger.Log("listening on "+service, 1)

	return nil
}
<|endoftext|>