<commit_before>//
// postgres db backend
//
package srnd

import (
  "database/sql"
  "fmt"
  "log"
  "sort"
  "strings"
  "time"
  _ "github.com/lib/pq"
)

type PostgresDatabase struct {
  conn *sql.DB
  db_str string
}

func NewPostgresDatabase(host, port, user, password string) Database {
  var db PostgresDatabase
  db.db_str = fmt.Sprintf("user=%s password=%s host=%s port=%s", user, password, host, port)
  db.conn = db.Conn()
  return db
}

func (self PostgresDatabase) Login() string {
  return self.db_str
}

// finalize all transactions
// close database connections
func (self PostgresDatabase) Close() {
  self.Conn().Close()
}

func (self PostgresDatabase) Conn() *sql.DB {
  if self.conn == nil {
    log.Println("connecting to database")
    var err error
    self.conn, err = sql.Open("postgres", self.Login())
    if err != nil {
      log.Fatalf("cannot open connection to db: %s", err)
    }
  }
  return self.conn
}

// create all tables
// will panic on fail
func (self PostgresDatabase) CreateTables() {
  tables := make(map[string]string)

  // table of active newsgroups
  tables["Newsgroups"] = `(
                            name VARCHAR(255) PRIMARY KEY,
                            last_post INTEGER NOT NULL,
                            restricted BOOLEAN
                          )`

  // table for storing nntp article meta data
  tables["Articles"] = `(
                          message_id VARCHAR(255) PRIMARY KEY,
                          message_id_hash VARCHAR(40) UNIQUE NOT NULL,
                          message_newsgroup VARCHAR(255),
                          message_ref_id VARCHAR(255),
                          time_obtained INTEGER NOT NULL,
                          FOREIGN KEY(message_newsgroup) REFERENCES Newsgroups(name)
                        )`

  // table for storing nntp article post content
  tables["ArticlePosts"] = `(
                              newsgroup VARCHAR(255),
                              message_id VARCHAR(255),
                              ref_id VARCHAR(255),
                              name TEXT NOT NULL,
                              subject TEXT NOT NULL,
                              path TEXT NOT NULL,
                              time_posted INTEGER NOT NULL,
                              message TEXT NOT NULL
                            )`

  // table for storing nntp article attachment info
  tables["ArticleAttachments"] = `(
                                    message_id VARCHAR(255),
                                    sha_hash VARCHAR(128) NOT NULL,
                                    filename TEXT NOT NULL,
                                    filepath TEXT NOT NULL
                                  )`

  for k, v := range(tables) {
    _, err := self.Conn().Exec(fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s%s", k, v))
    if err != nil {
      log.Fatalf("cannot create table %s, %s, login was '%s'", k, err,self.db_str)
    }
  }
}

func (self PostgresDatabase) GetRootPostsForExpiration(newsgroup string, threadcount int) []string {

  //TODO, do this all in 1 query with no bullshit after logic

  // root post -> last bump
  threads := make(map[string]int64)
  var rows *sql.Rows
  // get all posts for this newsgroup sorted by oldest post first
  stmt, err := self.Conn().Prepare("SELECT message_id, time_posted, subject, ref_id FROM ArticlePosts WHERE newsgroup = $1 ORDER BY time_posted ASC")
  if err != nil {
    log.Println("failed to prepare query for post expiration step 1", err)
    return nil
  }
  defer stmt.Close()
  rows, err = stmt.Query(newsgroup)
  if err != nil {
    log.Println("failed to execute query for post expiration step 1", err)
    return nil
  }
  // get results
  for rows.Next() {
    var msgid, subject, ref string
    var posted int64
    rows.Scan(&msgid, &posted, &subject, &ref)
    // is this a root post ?
    if len(ref) == 0 {
      // ya
      // record it bumped
      threads[msgid] = posted
      continue
    }
    // check for sage
    subject = strings.ToLower(subject)
    if strings.HasPrefix(subject, "sage ") || subject == "sage" {
      // this is a sage
      // we won't add it to the bump stuff
      continue
    }
    // bump the thread if the root post is there
    bump, ok := threads[ref]
    if ok {
      // bump it if it needs to
      if bump < posted {
        threads[ref] = posted
      }
    }
  }

  // make map such that: last bumped -> root post
  threads_out := make(map[int64]string)
  var bumps int64Sorter
  for root, bump := range threads {
    threads_out[bump] = root
    bumps = append(bumps, bump)
  }
  //sort by oldest first
  sort.Sort(bumps)
  var roots []string
  // add the oldest thread to the list of expired roots until we have enough threads left
  for len(bumps) >= threadcount {
    roots = append(roots, threads_out[bumps[0]])
    bumps = bumps[1:]
  }
  // return the list of expired roots
  return roots
}

func (self PostgresDatabase) GetAllNewsgroups() []string {
  var rows *sql.Rows
  var err error
  stmt, err := self.Conn().Prepare("SELECT name FROM Newsgroups")
  if err == nil {
    defer stmt.Close()
    rows, err = stmt.Query()
  }
  if err != nil {
    log.Println("failed to get all newsgroups", err)
    return nil
  }

  var groups []string

  if rows != nil {
    for rows.Next() {
      var group string
      rows.Scan(&group)
      groups = append(groups, group)
    }
  }
  return groups
}

func (self PostgresDatabase) GetThreadReplies(rootpost string, limit int) []string {
  var rows *sql.Rows
  var err error
  if limit > 0 {
    stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE message_id IN ( SELECT message_id FROM ArticlesPosts WHERE ref_id = $1 ORDER BY time_posted DESC LIMIT $2 ) ORDER BY time_posted ASC")
    if err == nil {
      defer stmt.Close()
      rows, err = stmt.Query(rootpost, limit)
    }
  } else {
    stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE message_id IN ( SELECT message_id FROM ArticlePosts WHERE ref_id = $1 ) ORDER BY time_posted ASC")
    if err == nil {
      defer stmt.Close()
      rows, err = stmt.Query(rootpost)
    }

  }
  if err != nil {
    log.Println("failed to get thread replies", rootpost, err)
    return nil
  }

  if rows == nil {
    return nil
  }

  var repls []string

  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    repls = append(repls, msgid)
  }
  log.Printf("thread has %d replies", len(repls))
  return repls

}

func (self PostgresDatabase) ThreadHasReplies(rootpost string) bool {
  log.Println("checking for replies for", rootpost)
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts WHERE ref_id = $1")
  if err != nil {
    log.Println("failed to prepare query to check for thread replies", rootpost, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(rootpost).Scan(&count)
  log.Printf("we have %d replies", count)
  return count > 0
}

func (self PostgresDatabase) GetGroupThreads(group string, recv chan string) {
  stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE newsgroup = $1 AND ref_id = '' ")
  if err != nil {
    log.Println("failed to prepare query to check for board threads", group, err)
    return
  }
  defer stmt.Close()
  rows, err := stmt.Query(group)
  if err != nil {
    log.Println("failed to execute query to check for board threads", group, err)
  }
  defer rows.Close()
  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    recv <- msgid
  }
}

func (self PostgresDatabase) GetLastBumpedThreads(threads int) []string {
  // TODO: detect sage
  stmt, err := self.Conn().Prepare("SELECT message_id, ref_id, time_posted FROM ArticlePosts ORDER BY time_posted DESC")
  if err != nil {
    log.Println("failed to prepare query for get last bumped", err)
    return nil
  }
  defer stmt.Close()
  rows, err := stmt.Query(threads)
  if err != nil {
    log.Println("failed to execute query for get last bumped", err)
  }
  defer rows.Close()

  var roots []string
  for rows.Next() {
    var msgid, refid string
    var posted int64
    rows.Scan(&msgid, &refid, &posted)
    if refid != "" {
      msgid = refid
    }
    put := true
    if len(roots) > 0 {
      for _, root := range roots {
        if root == msgid {
          put = false
          break
        }
      }
    }
    if put {
      roots = append(roots, msgid)
    }
    if len(roots) == threads {
      break
    }
  }
  return roots
}

func (self PostgresDatabase) GroupHasPosts(group string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts WHERE newsgroup = $1")
  if err != nil {
    log.Println("failed to prepare query to check for newsgroup posts", group, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(group).Scan(&count)
  return count > 0
}

// check if a newsgroup exists
func (self PostgresDatabase) HasNewsgroup(group string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(name) FROM Newsgroups WHERE name = $1")
  if err != nil {
    log.Println("failed to prepare query to check for newsgroup", group, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(group).Scan(&count)
  return count > 0
}

// check if an article exists
func (self PostgresDatabase) HasArticle(message_id string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM Articles WHERE message_id = $1")
  if err != nil {
    log.Println("failed to prepare query to check for article", message_id, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(message_id).Scan(&count)
  return count > 0
}

// check if an article exists
func (self PostgresDatabase) ArticleCount() int64 {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts")
  if err != nil {
    log.Println("failed to prepare query to get article count", err)
    return -1
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow().Scan(&count)
  return count
}

// register a new newsgroup
func (self PostgresDatabase) RegisterNewsgroup(group string) {
  stmt, err := self.Conn().Prepare("INSERT INTO Newsgroups (name, last_post) VALUES($1, $2)")
  if err != nil {
    log.Println("failed to prepare query to register newsgroup", group, err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(group, time.Now().Unix())
  if err != nil {
    log.Println("failed to register newsgroup", err)
  }
}

func (self PostgresDatabase) GetPostAttachments(messageID string) []string {
  var atts []string
  stmt, err := self.Conn().Prepare("SELECT filepath FROM ArticleAttachments WHERE message_id = $1")
  if err != nil {
    log.Println("failed to prepare query to get attachments for ", messageID, err)
    return atts
  }
  defer stmt.Close()
  rows, err := stmt.Query(messageID)
  if err != nil {
    log.Println("failed to execute query to get attachments for ", messageID, err)
  }
  defer rows.Close()
  for rows.Next() {
    var val string
    rows.Scan(&val)
    atts = append(atts, val)
  }
  return atts
}

// register a message with the database
func (self PostgresDatabase) RegisterArticle(message *NNTPMessage) {
  if ! self.HasNewsgroup(message.Newsgroup) {
    self.RegisterNewsgroup(message.Newsgroup)
  }
  // insert article metadata
  stmt, err := self.Conn().Prepare("INSERT INTO Articles (message_id, message_id_hash, message_newsgroup, time_obtained, message_ref_id) VALUES($1, $2, $3, $4, $5)")
  if err != nil {
    log.Println("failed to prepare query to register article", message.MessageID, err)
    return
  }
  defer stmt.Close()
  now := time.Now().Unix()
  _, err = stmt.Exec(message.MessageID, HashMessageID(message.MessageID), message.Newsgroup, now, message.Reference)
  if err != nil {
    log.Println("failed to register article", err)
  }
  // update newsgroup
  stmt, err = self.Conn().Prepare("UPDATE Newsgroups SET last_post = $1 WHERE name = $2")
  if err != nil {
    log.Println("cannot prepare query to update newsgroup last post", err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(now, message.Newsgroup)
  if err != nil {
    log.Println("cannot execute query to update newsgroup last post", err)
    return
  }
  // insert article post
  stmt, err = self.Conn().Prepare("INSERT INTO ArticlePosts(newsgroup, message_id, ref_id, name, subject, path, time_posted, message) VALUES($1, $2, $3, $4, $5, $6, $7, $8)")
  if err != nil {
    log.Println("cannot prepare query to insert article post", err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(message.Newsgroup, message.MessageID, message.Reference, message.Name, message.Subject, message.Path, message.Posted, message.Message)
  if err != nil {
    log.Println("cannot insert article post", err)
    return
  }
  // register all attachments
  for _, att := range message.Attachments {
    stmt, err = self.Conn().Prepare("INSERT INTO ArticleAttachments(message_id, sha_hash, filename, filepath) VALUES($1, $2, $3, $4)")
    if err != nil {
      log.Println("failed to prepare query to register attachment", err)
    }
    defer stmt.Close()
    _, err = stmt.Exec(message.MessageID, att.Hash(), att.Name, att.Filename())
    if err != nil {
      log.Println("failed to execute query to register attachment", err)
    }
  }
}

// get all articles in a newsgroup
// send result down a channel
func (self PostgresDatabase) GetAllArticlesInGroup(group string, recv chan string) {
  stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE newsgroup = $1")
  if err != nil {
    log.Printf("failed to prepare query for getting all articles in %s: %s", group, err)
    return
  }
  defer stmt.Close()
  rows, err := stmt.Query(group)
  if err != nil {
    log.Printf("Failed to execute quert for getting all articles in %s: %s", group, err)
    return
  }
  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    recv <- msgid
  }
}

// get all articles
// send result down a channel
func (self PostgresDatabase) GetAllArticles() []ArticleEntry {
  stmt, err := self.Conn().Prepare("SELECT message_id, newsgroup FROM ArticlePosts")
  if err != nil {
    log.Printf("failed to prepare query for getting all articles: %s", err)
    return nil
  }
  defer stmt.Close()
  rows, err := stmt.Query()
  if err != nil {
    log.Printf("Failed to execute quert for getting all articles: %s", err)
    return nil
  }
  var articles []ArticleEntry
  for rows.Next() {
    var entry ArticleEntry
    rows.Scan(&entry[0], &entry[1])
    articles = append(articles, entry)
  }
  return articles
}
<commit_msg>fix query<commit_after>//
// postgres db backend
//
package srnd

import (
  "database/sql"
  "fmt"
  "log"
  "sort"
  "strings"
  "time"
  _ "github.com/lib/pq"
)

type PostgresDatabase struct {
  conn *sql.DB
  db_str string
}

func NewPostgresDatabase(host, port, user, password string) Database {
  var db PostgresDatabase
  db.db_str = fmt.Sprintf("user=%s password=%s host=%s port=%s", user, password, host, port)
  db.conn = db.Conn()
  return db
}

func (self PostgresDatabase) Login() string {
  return self.db_str
}

// finalize all transactions
// close database connections
func (self PostgresDatabase) Close() {
  self.Conn().Close()
}

func (self PostgresDatabase) Conn() *sql.DB {
  if self.conn == nil {
    log.Println("connecting to database")
    var err error
    self.conn, err = sql.Open("postgres", self.Login())
    if err != nil {
      log.Fatalf("cannot open connection to db: %s", err)
    }
  }
  return self.conn
}

// create all tables
// will panic on fail
func (self PostgresDatabase) CreateTables() {
  tables := make(map[string]string)

  // table of active newsgroups
  tables["Newsgroups"] = `(
                            name VARCHAR(255) PRIMARY KEY,
                            last_post INTEGER NOT NULL,
                            restricted BOOLEAN
                          )`

  // table for storing nntp article meta data
  tables["Articles"] = `(
                          message_id VARCHAR(255) PRIMARY KEY,
                          message_id_hash VARCHAR(40) UNIQUE NOT NULL,
                          message_newsgroup VARCHAR(255),
                          message_ref_id VARCHAR(255),
                          time_obtained INTEGER NOT NULL,
                          FOREIGN KEY(message_newsgroup) REFERENCES Newsgroups(name)
                        )`

  // table for storing nntp article post content
  tables["ArticlePosts"] = `(
                              newsgroup VARCHAR(255),
                              message_id VARCHAR(255),
                              ref_id VARCHAR(255),
                              name TEXT NOT NULL,
                              subject TEXT NOT NULL,
                              path TEXT NOT NULL,
                              time_posted INTEGER NOT NULL,
                              message TEXT NOT NULL
                            )`

  // table for storing nntp article attachment info
  tables["ArticleAttachments"] = `(
                                    message_id VARCHAR(255),
                                    sha_hash VARCHAR(128) NOT NULL,
                                    filename TEXT NOT NULL,
                                    filepath TEXT NOT NULL
                                  )`

  for k, v := range(tables) {
    _, err := self.Conn().Exec(fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s%s", k, v))
    if err != nil {
      log.Fatalf("cannot create table %s, %s, login was '%s'", k, err,self.db_str)
    }
  }
}

func (self PostgresDatabase) GetRootPostsForExpiration(newsgroup string, threadcount int) []string {

  //TODO, do this all in 1 query with no bullshit after logic

  // root post -> last bump
  threads := make(map[string]int64)
  var rows *sql.Rows
  // get all posts for this newsgroup sorted by oldest post first
  stmt, err := self.Conn().Prepare("SELECT message_id, time_posted, subject, ref_id FROM ArticlePosts WHERE newsgroup = $1 ORDER BY time_posted ASC")
  if err != nil {
    log.Println("failed to prepare query for post expiration step 1", err)
    return nil
  }
  defer stmt.Close()
  rows, err = stmt.Query(newsgroup)
  if err != nil {
    log.Println("failed to execute query for post expiration step 1", err)
    return nil
  }
  // get results
  for rows.Next() {
    var msgid, subject, ref string
    var posted int64
    rows.Scan(&msgid, &posted, &subject, &ref)
    // is this a root post ?
    if len(ref) == 0 {
      // ya
      // record it bumped
      threads[msgid] = posted
      continue
    }
    // check for sage
    subject = strings.ToLower(subject)
    if strings.HasPrefix(subject, "sage ") || subject == "sage" {
      // this is a sage
      // we won't add it to the bump stuff
      continue
    }
    // bump the thread if the root post is there
    bump, ok := threads[ref]
    if ok {
      // bump it if it needs to
      if bump < posted {
        threads[ref] = posted
      }
    }
  }

  // make map such that: last bumped -> root post
  threads_out := make(map[int64]string)
  var bumps int64Sorter
  for root, bump := range threads {
    threads_out[bump] = root
    bumps = append(bumps, bump)
  }
  //sort by oldest first
  sort.Sort(bumps)
  var roots []string
  // add the oldest thread to the list of expired roots until we have enough threads left
  for len(bumps) >= threadcount {
    roots = append(roots, threads_out[bumps[0]])
    bumps = bumps[1:]
  }
  // return the list of expired roots
  return roots
}

func (self PostgresDatabase) GetAllNewsgroups() []string {
  var rows *sql.Rows
  var err error
  stmt, err := self.Conn().Prepare("SELECT name FROM Newsgroups")
  if err == nil {
    defer stmt.Close()
    rows, err = stmt.Query()
  }
  if err != nil {
    log.Println("failed to get all newsgroups", err)
    return nil
  }

  var groups []string

  if rows != nil {
    for rows.Next() {
      var group string
      rows.Scan(&group)
      groups = append(groups, group)
    }
  }
  return groups
}

func (self PostgresDatabase) GetThreadReplies(rootpost string, limit int) []string {
  var rows *sql.Rows
  var err error
  if limit > 0 {
    stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE message_id IN ( SELECT message_id FROM ArticlesPosts WHERE ref_id = $1 ORDER BY time_posted DESC LIMIT $2 ) ORDER BY time_posted ASC")
    if err == nil {
      defer stmt.Close()
      rows, err = stmt.Query(rootpost, limit)
    }
  } else {
    stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE message_id IN ( SELECT message_id FROM ArticlePosts WHERE ref_id = $1 ) ORDER BY time_posted ASC")
    if err == nil {
      defer stmt.Close()
      rows, err = stmt.Query(rootpost)
    }

  }
  if err != nil {
    log.Println("failed to get thread replies", rootpost, err)
    return nil
  }

  if rows == nil {
    return nil
  }

  var repls []string

  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    repls = append(repls, msgid)
  }
  log.Printf("thread has %d replies", len(repls))
  return repls

}

func (self PostgresDatabase) ThreadHasReplies(rootpost string) bool {
  log.Println("checking for replies for", rootpost)
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts WHERE ref_id = $1")
  if err != nil {
    log.Println("failed to prepare query to check for thread replies", rootpost, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(rootpost).Scan(&count)
  log.Printf("we have %d replies", count)
  return count > 0
}

func (self PostgresDatabase) GetGroupThreads(group string, recv chan string) {
  stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE newsgroup = $1 AND ref_id = '' ")
  if err != nil {
    log.Println("failed to prepare query to check for board threads", group, err)
    return
  }
  defer stmt.Close()
  rows, err := stmt.Query(group)
  if err != nil {
    log.Println("failed to execute query to check for board threads", group, err)
  }
  defer rows.Close()
  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    recv <- msgid
  }
}

func (self PostgresDatabase) GetLastBumpedThreads(threads int) []string {
  // TODO: detect sage
  stmt, err := self.Conn().Prepare("SELECT message_id, ref_id, time_posted FROM ArticlePosts ORDER BY time_posted DESC")
  if err != nil {
    log.Println("failed to prepare query for get last bumped", err)
    return nil
  }
  defer stmt.Close()
  rows, err := stmt.Query()
  if err != nil {
    log.Println("failed to execute query for get last bumped", err)
  }
  defer rows.Close()

  var roots []string
  for rows.Next() {
    var msgid, refid string
    var posted int64
    rows.Scan(&msgid, &refid, &posted)
    if refid != "" {
      msgid = refid
    }
    put := true
    if len(roots) > 0 {
      for _, root := range roots {
        if root == msgid {
          put = false
          break
        }
      }
    }
    if put {
      roots = append(roots, msgid)
    }
    if len(roots) == threads {
      break
    }
  }
  return roots
}

func (self PostgresDatabase) GroupHasPosts(group string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts WHERE newsgroup = $1")
  if err != nil {
    log.Println("failed to prepare query to check for newsgroup posts", group, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(group).Scan(&count)
  return count > 0
}

// check if a newsgroup exists
func (self PostgresDatabase) HasNewsgroup(group string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(name) FROM Newsgroups WHERE name = $1")
  if err != nil {
    log.Println("failed to prepare query to check for newsgroup", group, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(group).Scan(&count)
  return count > 0
}

// check if an article exists
func (self PostgresDatabase) HasArticle(message_id string) bool {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM Articles WHERE message_id = $1")
  if err != nil {
    log.Println("failed to prepare query to check for article", message_id, err)
    return false
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow(message_id).Scan(&count)
  return count > 0
}

// check if an article exists
func (self PostgresDatabase) ArticleCount() int64 {
  stmt, err := self.Conn().Prepare("SELECT COUNT(message_id) FROM ArticlePosts")
  if err != nil {
    log.Println("failed to prepare query to get article count", err)
    return -1
  }
  defer stmt.Close()
  var count int64
  stmt.QueryRow().Scan(&count)
  return count
}

// register a new newsgroup
func (self PostgresDatabase) RegisterNewsgroup(group string) {
  stmt, err := self.Conn().Prepare("INSERT INTO Newsgroups (name, last_post) VALUES($1, $2)")
  if err != nil {
    log.Println("failed to prepare query to register newsgroup", group, err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(group, time.Now().Unix())
  if err != nil {
    log.Println("failed to register newsgroup", err)
  }
}

func (self PostgresDatabase) GetPostAttachments(messageID string) []string {
  var atts []string
  stmt, err := self.Conn().Prepare("SELECT filepath FROM ArticleAttachments WHERE message_id = $1")
  if err != nil {
    log.Println("failed to prepare query to get attachments for ", messageID, err)
    return atts
  }
  defer stmt.Close()
  rows, err := stmt.Query(messageID)
  if err != nil {
    log.Println("failed to execute query to get attachments for ", messageID, err)
  }
  defer rows.Close()
  for rows.Next() {
    var val string
    rows.Scan(&val)
    atts = append(atts, val)
  }
  return atts
}

// register a message with the database
func (self PostgresDatabase) RegisterArticle(message *NNTPMessage) {
  if ! self.HasNewsgroup(message.Newsgroup) {
    self.RegisterNewsgroup(message.Newsgroup)
  }
  // insert article metadata
  stmt, err := self.Conn().Prepare("INSERT INTO Articles (message_id, message_id_hash, message_newsgroup, time_obtained, message_ref_id) VALUES($1, $2, $3, $4, $5)")
  if err != nil {
    log.Println("failed to prepare query to register article", message.MessageID, err)
    return
  }
  defer stmt.Close()
  now := time.Now().Unix()
  _, err = stmt.Exec(message.MessageID, HashMessageID(message.MessageID), message.Newsgroup, now, message.Reference)
  if err != nil {
    log.Println("failed to register article", err)
  }
  // update newsgroup
  stmt, err = self.Conn().Prepare("UPDATE Newsgroups SET last_post = $1 WHERE name = $2")
  if err != nil {
    log.Println("cannot prepare query to update newsgroup last post", err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(now, message.Newsgroup)
  if err != nil {
    log.Println("cannot execute query to update newsgroup last post", err)
    return
  }
  // insert article post
  stmt, err = self.Conn().Prepare("INSERT INTO ArticlePosts(newsgroup, message_id, ref_id, name, subject, path, time_posted, message) VALUES($1, $2, $3, $4, $5, $6, $7, $8)")
  if err != nil {
    log.Println("cannot prepare query to insert article post", err)
    return
  }
  defer stmt.Close()
  _, err = stmt.Exec(message.Newsgroup, message.MessageID, message.Reference, message.Name, message.Subject, message.Path, message.Posted, message.Message)
  if err != nil {
    log.Println("cannot insert article post", err)
    return
  }
  // register all attachments
  for _, att := range message.Attachments {
    stmt, err = self.Conn().Prepare("INSERT INTO ArticleAttachments(message_id, sha_hash, filename, filepath) VALUES($1, $2, $3, $4)")
    if err != nil {
      log.Println("failed to prepare query to register attachment", err)
    }
    defer stmt.Close()
    _, err = stmt.Exec(message.MessageID, att.Hash(), att.Name, att.Filename())
    if err != nil {
      log.Println("failed to execute query to register attachment", err)
    }
  }
}

// get all articles in a newsgroup
// send result down a channel
func (self PostgresDatabase) GetAllArticlesInGroup(group string, recv chan string) {
  stmt, err := self.Conn().Prepare("SELECT message_id FROM ArticlePosts WHERE newsgroup = $1")
  if err != nil {
    log.Printf("failed to prepare query for getting all articles in %s: %s", group, err)
    return
  }
  defer stmt.Close()
  rows, err := stmt.Query(group)
  if err != nil {
    log.Printf("Failed to execute quert for getting all articles in %s: %s", group, err)
    return
  }
  for rows.Next() {
    var msgid string
    rows.Scan(&msgid)
    recv <- msgid
  }
}

// get all articles
// send result down a channel
func (self PostgresDatabase) GetAllArticles() []ArticleEntry {
  stmt, err := self.Conn().Prepare("SELECT message_id, newsgroup FROM ArticlePosts")
  if err != nil {
    log.Printf("failed to prepare query for getting all articles: %s", err)
    return nil
  }
  defer stmt.Close()
  rows, err := stmt.Query()
  if err != nil {
    log.Printf("Failed to execute quert for getting all articles: %s", err)
    return nil
  }
  var articles []ArticleEntry
  for rows.Next() {
    var entry ArticleEntry
    rows.Scan(&entry[0], &entry[1])
    articles = append(articles, entry)
  }
  return articles
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Damon Revoe. All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package main

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v2"
)

type workspaceParams struct {
	PkgPath    string `yaml:"pkgpath"`
	InstallDir string `yaml:"installdir,omitempty"`
}

func getWorkspaceDir() string {
	workspaceDir := "." + appName

	if flags.workspaceDir != "" {
		workspaceDir = filepath.Join(flags.workspaceDir, workspaceDir)
	}

	return workspaceDir
}

func getInitPathname(workspaceDir string) string {
	return filepath.Join(workspaceDir, "init.yaml")
}

func createWorkspace() (*workspaceParams, error) {
	workspaceDir := getWorkspaceDir()

	if _, err := os.Stat(workspaceDir); err == nil {
		return nil, errors.New("workspace already initialized")
	}

	pkgpath, err := getPackagePathFromEnvironment()

	if err != nil {
		return nil, err
	}

	wp := workspaceParams{pkgpath, flags.installDir}

	out, err := yaml.Marshal(&wp)
	if err != nil {
		return nil, err
	}

	err = os.Mkdir(workspaceDir, os.FileMode(0775))
	if err != nil {
		return nil, err
	}

	err = ioutil.WriteFile(getInitPathname(workspaceDir),
		out, os.FileMode(0664))
	if err != nil {
		return nil, err
	}

	return &workspaceParams{}, nil
}

func readWorkspaceParams() (*workspaceParams, error) {
	in, err := ioutil.ReadFile(getInitPathname(getWorkspaceDir()))
	if err != nil {
		return nil, err
	}

	var wp workspaceParams
	err = yaml.Unmarshal(in, &wp)

	return &wp, err
}
<commit_msg>Rename init.yaml to settings.yaml<commit_after>// Copyright (C) 2017 Damon Revoe. All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package main

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v2"
)

type workspaceParams struct {
	PkgPath    string `yaml:"pkgpath"`
	InstallDir string `yaml:"installdir,omitempty"`
}

func getWorkspaceDir() string {
	workspaceDir := "." + appName

	if flags.workspaceDir != "" {
		workspaceDir = filepath.Join(flags.workspaceDir, workspaceDir)
	}

	return workspaceDir
}

func getPathToSettings(workspaceDir string) string {
	return filepath.Join(workspaceDir, "settings.yaml")
}

func createWorkspace() (*workspaceParams, error) {
	workspaceDir := getWorkspaceDir()

	if _, err := os.Stat(workspaceDir); err == nil {
		return nil, errors.New("workspace already initialized")
	}

	pkgpath, err := getPackagePathFromEnvironment()

	if err != nil {
		return nil, err
	}

	wp := workspaceParams{pkgpath, flags.installDir}

	out, err := yaml.Marshal(&wp)
	if err != nil {
		return nil, err
	}

	err = os.Mkdir(workspaceDir, os.FileMode(0775))
	if err != nil {
		return nil, err
	}

	err = ioutil.WriteFile(getPathToSettings(workspaceDir),
		out, os.FileMode(0664))
	if err != nil {
		return nil, err
	}

	return &workspaceParams{}, nil
}

func readWorkspaceParams() (*workspaceParams, error) {
	in, err := ioutil.ReadFile(getPathToSettings(getWorkspaceDir()))
	if err != nil {
		return nil, err
	}

	var wp workspaceParams
	err = yaml.Unmarshal(in, &wp)

	return &wp, err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"log"

	"github.com/dmnlk/gomadare"
	"github.com/dmnlk/stringUtils"
	"github.com/rem7/goprowl"
	"strings"
)

var (
	CONSUMER_KEY        string
	CONSUMER_KEY_SECRET string
	ACCESS_TOKEN        string
	ACCESS_TOKEN_SECRET string
	PROWL_API_KEY       string
	PROWL               goprowl.Goprowl
)

func main() {
	err := configureToken()
	if err != nil {
		fmt.Println(err)
		return
	}
	err = PROWL.RegisterKey(PROWL_API_KEY)
	if err != nil {
		fmt.Println(err)
		return
	}

	client := gomadare.NewClient(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET)
	fmt.Println("aa")
	client.GetUserStream(nil, func(s gomadare.Status, e gomadare.Event) {
		if &s != nil {
			go sendReplyAndRetweetToProwl(s)
		}
		if &e != nil {
			go sendEventToProwl(e)
		}
	})
}

func configureToken() error {
	CONSUMER_KEY = os.Getenv("CONSUMER_KEY")
	CONSUMER_KEY_SECRET = os.Getenv("CONSUMER_KEY_SECRET")
	ACCESS_TOKEN = os.Getenv("ACCESS_TOKEN")
	ACCESS_TOKEN_SECRET = os.Getenv("ACCESS_TOKEN_SECRET")
	PROWL_API_KEY = os.Getenv("PROWL_API_KEY")
	if ng := stringUtils.IsAnyEmpty(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET, PROWL_API_KEY); ng {
		return fmt.Errorf("some key invalid")
	}

	return nil
}

func sendEventToProwl(e gomadare.Event) {
	if stringUtils.IsEmpty(e.Event) {
		return
	}
	emoji := getEventEmoji(e)
	n := &goprowl.Notification{
		Application: "yggdrasill",
		Description: emoji + " " + e.TargetObject.Text,
		Event:       e.Event + " by " + e.Source.ScreenName,
		Priority:    "1",
	}

	PROWL.Push(n)
}

func getEventEmoji(event gomadare.Event) string {
	if event.Event == "favorite" {
		return "\u2b50"
	}
	if event.Event == "unfavorite" {
		return "\U0001f44e"
	}
	if event.Event == "list_member_removed" {
		return "\u274c"
	}
	if event.Event == "list_member_added" {
		return "\u2755"
	}
	if event.Event == "follow" {
		return "\u2661"
	}
	log.Println("unknown event:" + event.Event)
	return event.Event
}

func getProwlNotification(event gomadare.Event) goprowl.Notification {
	n := &goprowl.Notification{
		Application: "Yggdrsill",
		Priority: "1",
	}

	return *n
}

func sendReplyAndRetweetToProwl(s gomadare.Status) {
	// reply Event
	if len(s.Entities.UserMentions) > 0 {
		for _, mention := range s.Entities.UserMentions {
			if mention.ScreenName == "dmnlk" {
				var n *goprowl.Notification
				if strings.Contains(s.Text, "RT") {
					n = &goprowl.Notification{
						Application: "Golang",
						Description: "\U0001f4a1" + " " + s.Text,
						Event:       "RT by " + s.User.ScreenName,
						Priority:    "1",
					}
				} else {
					n = &goprowl.Notification{
						Application: "Golang",
						Description: "\U0001f4a1" + " " + s.Text,
						Event:       "Mentioned by " + s.User.ScreenName,
						Priority:    "1",
					}
				}

				PROWL.Push(n)
			}
		}
	}

	//RTイベント
	if len(s.RetweetedStatus) > 0 {
		n := &goprowl.Notification{
			Application: "Golang",
			Description: "\U0001f4a1" + " " + s.Text,
			Event:       "RT by " + s.User.ScreenName,
			Priority:    "1",
		}

		PROWL.Push(n)
	}
}
<commit_msg>fix retweet check<commit_after>package main

import (
	"fmt"
	"os"

	"log"

	"github.com/dmnlk/gomadare"
	"github.com/dmnlk/stringUtils"
	"github.com/rem7/goprowl"
	"strings"
)

var (
	CONSUMER_KEY        string
	CONSUMER_KEY_SECRET string
	ACCESS_TOKEN        string
	ACCESS_TOKEN_SECRET string
	PROWL_API_KEY       string
	PROWL               goprowl.Goprowl
)

func main() {
	err := configureToken()
	if err != nil {
		fmt.Println(err)
		return
	}
	err = PROWL.RegisterKey(PROWL_API_KEY)
	if err != nil {
		fmt.Println(err)
		return
	}

	client := gomadare.NewClient(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET)

	client.GetUserStream(nil, func(s gomadare.Status, e gomadare.Event) {
		if &s != nil {
			go sendReplyAndRetweetToProwl(s)
		}
		if &e != nil {
			go sendEventToProwl(e)
		}
	})
}

func configureToken() error {
	CONSUMER_KEY = os.Getenv("CONSUMER_KEY")
	CONSUMER_KEY_SECRET = os.Getenv("CONSUMER_KEY_SECRET")
	ACCESS_TOKEN = os.Getenv("ACCESS_TOKEN")
	ACCESS_TOKEN_SECRET = os.Getenv("ACCESS_TOKEN_SECRET")
	PROWL_API_KEY = os.Getenv("PROWL_API_KEY")
	if ng := stringUtils.IsAnyEmpty(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET, PROWL_API_KEY); ng {
		return fmt.Errorf("some key invalid")
	}

	return nil
}

func sendEventToProwl(e gomadare.Event) {
	if stringUtils.IsEmpty(e.Event) {
		return
	}
	emoji := getEventEmoji(e)
	n := &goprowl.Notification{
		Application: "yggdrasill",
		Description: emoji + " " + e.TargetObject.Text,
		Event:       e.Event + " by " + e.Source.ScreenName,
		Priority:    "1",
	}

	PROWL.Push(n)
}

func getEventEmoji(event gomadare.Event) string {
	if event.Event == "favorite" {
		return "\u2b50"
	}
	if event.Event == "unfavorite" {
		return "\U0001f44e"
	}
	if event.Event == "list_member_removed" {
		return "\u274c"
	}
	if event.Event == "list_member_added" {
		return "\u2755"
	}
	if event.Event == "follow" {
		return "\u2661"
	}
	log.Println("unknown event:" + event.Event)
	return event.Event
}

func getProwlNotification(event gomadare.Event) goprowl.Notification {
	n := &goprowl.Notification{
		Application: "Yggdrsill",
		Priority: "1",
	}

	return *n
}

func sendReplyAndRetweetToProwl(s gomadare.Status) {
	// reply Event
	if len(s.Entities.UserMentions) > 0 {
		for _, mention := range s.Entities.UserMentions {
			if mention.ScreenName == "dmnlk" {
				var n *goprowl.Notification
				if strings.Contains(s.Text, "RT") {
					n = &goprowl.Notification{
						Application: "yggdrasill",
						Description: "\U0001f4a1" + " " + s.Text,
						Event:       "RT by " + s.User.ScreenName,
						Priority:    "1",
					}
				} else {
					n = &goprowl.Notification{
						Application: "yggdrasill",
						Description: "\U0001f4a1" + " " + s.Text,
						Event:       "Mentioned by " + s.User.ScreenName,
						Priority:    "1",
					}
				}

				PROWL.Push(n)
			}
		}
	}

	//RTイベント
	if s.RetweetedStatus.User.ScreenName == "dmnlk" {
		n := &goprowl.Notification{
			Application: "yggdrasill",
			Description: "\U0001f4a1" + " " + s.Text,
			Event:       "RT by " + s.User.ScreenName,
			Priority:    "1",
		}

		PROWL.Push(n)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

// Utility code for working with OAuth. Subject to interface change.
package oauthutil

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path"

	"golang.org/x/oauth2"
)

// Create an oauth2.TokenSource that retrieves tokens in the following
// preference order:
//
//  1. From a copy of the latest obtained token, if still valid.
//
//  2. From a cache on disk, when it contains a valid token. The supplied file
//     basename will be placed in the user's home directory.
//
//  3. From an oauth token refresher, seeded with the last valid cached token.
//
//  4. By exchanging a new authorization code obtained via the supplied flag.
//     If the flag has not been set, the program will be aborted with a helpful
//     message.
//
// Note in particular the warning about aborting the program in #4.
func NewTerribleTokenSource(
	config *oauth2.Config,
	authCodeFlag *flag.Flag,
	cacheFilename string) (oauth2.TokenSource, error) {
	// Expand the path to the token cache.
	homedir, err := getHomeDir()
	if err != nil {
		return nil, fmt.Errorf("Finding homedir: %v", err)
	}

	cachePath := path.Join(homedir, cacheFilename)

	// Create the underlying token source.
	var ts oauth2.TokenSource = &tokenSource{
		oauth2Config: config,
		authCodeFlag: authCodeFlag,
		cachePath:    cachePath,
	}

	// Make sure not to consult the cache when a valid token is already lying
	// around in memory.
	ts = oauth2.ReuseTokenSource(nil, ts)

	return ts, nil
}

func getHomeDir() (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", fmt.Errorf("user.Current: %v", err)
	}

	return usr.HomeDir, nil
}

type tokenSource struct {
	oauth2Config *oauth2.Config
	authCodeFlag *flag.Flag
	cachePath    string
}

func (ts *tokenSource) getAuthCode() string {
	s := ts.authCodeFlag.Value.String()
	if s == "" {
		// NOTE(jacobsa): As of 2015-02-05 the documentation for
		// oauth2.Config.AuthCodeURL says that it is required to set this, but as
		// far as I can tell (cf. RFC 6749 §10.12) it is irrelevant for an
		// installed application that doesn't have a meaningful redirect URL.
		const csrfToken = ""

		fmt.Printf("You must set -%s.\n", ts.authCodeFlag.Name)
		fmt.Println("Visit this URL to obtain a code:")
		fmt.Println("    ", ts.oauth2Config.AuthCodeURL(csrfToken, oauth2.AccessTypeOffline))
		os.Exit(1)
	}

	return s
}

func (ts *tokenSource) getToken() (*oauth2.Token, error) {
	// First consult the cache.
	log.Println("Looking up OAuth token in cache.")
	t, err := ts.LookUp()
	if err != nil {
		// Log the error and ignore it.
		log.Println("Error loading from token cache: ", err)
	}

	// If the cached token is valid, we can return it.
	if t != nil && t.Valid() {
		return t, nil
	}

	// Otherwise if there was a cached token, it is probably expired and we can
	// attempt to refresh it.
	if t != nil {
		if t, err = ts.oauth2Config.TokenSource(oauth2.NoContext, t).Token(); err == nil {
			return t, nil
		}

		log.Println("Error refreshing token:", err)
	}

	// We must fall back to exchanging an auth code.
	log.Println("Exchanging auth code for token.")
	return ts.oauth2Config.Exchange(oauth2.NoContext, ts.getAuthCode())
}

func (ts *tokenSource) Token() (*oauth2.Token, error) {
	// Obtain the token.
	t, err := ts.getToken()
	if err != nil {
		return nil, err
	}

	// Insert into cache, then return the token.
	err = ts.Insert(t)
	if err != nil {
		log.Println("Error inserting into token cache:", err)
	}

	return t, nil
}

// Look for a token in the cache. Returns nil, nil on miss.
func (ts *tokenSource) LookUp() (*oauth2.Token, error) {
	// Open the cache file.
	file, err := os.Open(ts.cachePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// Decode the token.
	t := &oauth2.Token{}
	if err := json.NewDecoder(file).Decode(t); err != nil {
		return nil, err
	}

	return t, nil
}

func (ts *tokenSource) Insert(t *oauth2.Token) error {
	const flags = os.O_RDWR | os.O_CREATE | os.O_TRUNC
	const perm = 0600

	// Open the cache file.
	file, err := os.OpenFile(ts.cachePath, flags, perm)
	if err != nil {
		return err
	}

	// Encode the token.
	if err := json.NewEncoder(file).Encode(t); err != nil {
		file.Close()
		return err
	}

	// Close the file.
	if err := file.Close(); err != nil {
		return err
	}

	return nil
}
<commit_msg>Catch nil flags.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

// Utility code for working with OAuth. Subject to interface change.
package oauthutil

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path"

	"golang.org/x/oauth2"
)

// Create an oauth2.TokenSource that retrieves tokens in the following
// preference order:
//
//  1. From a copy of the latest obtained token, if still valid.
//
//  2. From a cache on disk, when it contains a valid token. The supplied file
//     basename will be placed in the user's home directory.
//
//  3. From an oauth token refresher, seeded with the last valid cached token.
//
//  4. By exchanging a new authorization code obtained via the supplied flag.
//     If the flag has not been set, the program will be aborted with a helpful
//     message.
//
// Note in particular the warning about aborting the program in #4.
func NewTerribleTokenSource(
	config *oauth2.Config,
	authCodeFlag *flag.Flag,
	cacheFilename string) (oauth2.TokenSource, error) {
	// Catch the common mistake of an unknown flag.
	if authCodeFlag == nil {
		return nil, errors.New("NewTerribleTokenSource passed a nil flag.")
	}

	// Expand the path to the token cache.
	homedir, err := getHomeDir()
	if err != nil {
		return nil, fmt.Errorf("Finding homedir: %v", err)
	}

	cachePath := path.Join(homedir, cacheFilename)

	// Create the underlying token source.
	var ts oauth2.TokenSource = &tokenSource{
		oauth2Config: config,
		authCodeFlag: authCodeFlag,
		cachePath:    cachePath,
	}

	// Make sure not to consult the cache when a valid token is already lying
	// around in memory.
	ts = oauth2.ReuseTokenSource(nil, ts)

	return ts, nil
}

func getHomeDir() (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", fmt.Errorf("user.Current: %v", err)
	}

	return usr.HomeDir, nil
}

type tokenSource struct {
	oauth2Config *oauth2.Config
	authCodeFlag *flag.Flag
	cachePath    string
}

func (ts *tokenSource) getAuthCode() string {
	s := ts.authCodeFlag.Value.String()
	if s == "" {
		// NOTE(jacobsa): As of 2015-02-05 the documentation for
		// oauth2.Config.AuthCodeURL says that it is required to set this, but as
		// far as I can tell (cf. RFC 6749 §10.12) it is irrelevant for an
		// installed application that doesn't have a meaningful redirect URL.
		const csrfToken = ""

		fmt.Printf("You must set -%s.\n", ts.authCodeFlag.Name)
		fmt.Println("Visit this URL to obtain a code:")
		fmt.Println("    ", ts.oauth2Config.AuthCodeURL(csrfToken, oauth2.AccessTypeOffline))
		os.Exit(1)
	}

	return s
}

func (ts *tokenSource) getToken() (*oauth2.Token, error) {
	// First consult the cache.
	log.Println("Looking up OAuth token in cache.")
	t, err := ts.LookUp()
	if err != nil {
		// Log the error and ignore it.
		log.Println("Error loading from token cache: ", err)
	}

	// If the cached token is valid, we can return it.
	if t != nil && t.Valid() {
		return t, nil
	}

	// Otherwise if there was a cached token, it is probably expired and we can
	// attempt to refresh it.
	if t != nil {
		if t, err = ts.oauth2Config.TokenSource(oauth2.NoContext, t).Token(); err == nil {
			return t, nil
		}

		log.Println("Error refreshing token:", err)
	}

	// We must fall back to exchanging an auth code.
	log.Println("Exchanging auth code for token.")
	return ts.oauth2Config.Exchange(oauth2.NoContext, ts.getAuthCode())
}

func (ts *tokenSource) Token() (*oauth2.Token, error) {
	// Obtain the token.
	t, err := ts.getToken()
	if err != nil {
		return nil, err
	}

	// Insert into cache, then return the token.
	err = ts.Insert(t)
	if err != nil {
		log.Println("Error inserting into token cache:", err)
	}

	return t, nil
}

// Look for a token in the cache. Returns nil, nil on miss.
func (ts *tokenSource) LookUp() (*oauth2.Token, error) {
	// Open the cache file.
	file, err := os.Open(ts.cachePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// Decode the token.
	t := &oauth2.Token{}
	if err := json.NewDecoder(file).Decode(t); err != nil {
		return nil, err
	}

	return t, nil
}

func (ts *tokenSource) Insert(t *oauth2.Token) error {
	const flags = os.O_RDWR | os.O_CREATE | os.O_TRUNC
	const perm = 0600

	// Open the cache file.
	file, err := os.OpenFile(ts.cachePath, flags, perm)
	if err != nil {
		return err
	}

	// Encode the token.
	if err := json.NewEncoder(file).Encode(t); err != nil {
		file.Close()
		return err
	}

	// Close the file.
	if err := file.Close(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"

	"github.com/gorilla/websocket"
	"gopkg.in/readline.v1"
)

var (
	VERSION = "v0.2.1"

	showVersion = flag.Bool("version", false, "show version number and exit")
	verbose     = flag.Bool("v", false, "verbose output")
)

func main() {
	flag.Parse()

	if *showVersion || *verbose {
		fmt.Printf("ws-client %s \n", VERSION)
		fmt.Println("https://github.com/oliver006/ws-client/")
		fmt.Println()
		if *showVersion {
			os.Exit(0)
		}
	}

	if len(os.Args) < 2 {
		fmt.Println("Usage:")
		fmt.Println("  ws-client <<ws:// or wss:// URL>>")
		os.Exit(-1)
	}

	addr := os.Args[len(os.Args)-1]
	if !strings.HasPrefix(addr, "ws://") && !strings.HasPrefix(addr, "wss://") {
		addr = "ws://" + addr
	}

	if *verbose {
		fmt.Printf("connecting to %s\n", addr)
	}
	c, _, err := websocket.DefaultDialer.Dial(addr, nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
	defer c.Close()
	fmt.Printf("connected to  %s\n", addr)

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	rl, err := readline.NewEx(&readline.Config{
		Prompt: "\033[32m»\033[0m ",
	})
	if err != nil {
		panic(err)
	}

	defer rl.Close()

	go func() {
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				fmt.Printf("<<server: %s>>\n", err)
				close(interrupt)
				return
			}
			io.WriteString(rl.Stdout(), fmt.Sprintf("\033[31m«\033[0m %s\n", message))
		}
	}()

	go func() {
		for {
			line, err := rl.Readline()
			if err == readline.ErrInterrupt {
				interrupt <- os.Interrupt
				return
			} else if err != nil {
				close(interrupt)
				return
			}
			if len(line) == 0 {
				continue
			}
			err = c.WriteMessage(websocket.TextMessage, []byte(line))
			if err != nil {
				fmt.Println("err:", err)
				return
			}
		}
	}()

	select {
	case <-interrupt:
		c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		fmt.Println("<<client: sent websocket close frame>>")
		c.Close()
		os.Exit(0)
	}
}
<commit_msg>more info when not enough cmd line params<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"

	"github.com/gorilla/websocket"
	"gopkg.in/readline.v1"
)

var (
	VERSION = "v0.2.1"

	showVersion = flag.Bool("version", false, "show version number and exit")
	verbose     = flag.Bool("v", false, "verbose output")
)

func main() {
	flag.Parse()

	if *showVersion || *verbose || len(os.Args) < 2 {
		fmt.Printf("ws-client %s \n", VERSION)
		fmt.Println("https://github.com/oliver006/ws-client/")
		fmt.Println()
		if *showVersion {
			os.Exit(0)
		}
	}

	if len(os.Args) < 2 {
		fmt.Println("Usage:")
		fmt.Println("  ws-client <<ws:// or wss:// URL>>")
		os.Exit(-1)
	}

	addr := os.Args[len(os.Args)-1]
	if !strings.HasPrefix(addr, "ws://") && !strings.HasPrefix(addr, "wss://") {
		addr = "ws://" + addr
	}

	if *verbose {
		fmt.Printf("connecting to %s\n", addr)
	}
	c, _, err := websocket.DefaultDialer.Dial(addr, nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
	defer c.Close()
	fmt.Printf("connected to  %s\n", addr)

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	rl, err := readline.NewEx(&readline.Config{
		Prompt: "\033[32m»\033[0m ",
	})
	if err != nil {
		panic(err)
	}

	defer rl.Close()

	go func() {
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				fmt.Printf("<<server: %s>>\n", err)
				close(interrupt)
				return
			}
			io.WriteString(rl.Stdout(), fmt.Sprintf("\033[31m«\033[0m %s\n", message))
		}
	}()

	go func() {
		for {
			line, err := rl.Readline()
			if err == readline.ErrInterrupt {
				interrupt <- os.Interrupt
				return
			} else if err != nil {
				close(interrupt)
				return
			}
			if len(line) == 0 {
				continue
			}
			err = c.WriteMessage(websocket.TextMessage, []byte(line))
			if err != nil {
				fmt.Println("err:", err)
				return
			}
		}
	}()

	select {
	case <-interrupt:
		c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		fmt.Println("<<client: sent websocket close frame>>")
		c.Close()
		os.Exit(0)
	}
}
<|endoftext|>


<commit_before>package ws

import (
	"net/http"

	"github.com/gorilla/websocket"
	"github.com/marcuswestin/go-errs"
)

var (
	/*
		HTTP Upgrade Options

		These are the different options
	*/

	// CheckOrigin will be called by the upgrader on incoming requests to check
	// the origin. If the CheckOrigin function returns false, then the Upgrade
	// method fails the WebSocket handshake with HTTP status 403.
	// If nil, then use a safe default: fail the handshake if the Origin request
	// header is present and not equal to the Host request header.
	CheckOrigin func(r *http.Request) bool

	// ErrorFn specifies the function for generating HTTP error responses.
	// If nil, then http.Error is used to generate the HTTP response.
	ErrorFn func(w http.ResponseWriter, r *http.Request, status int, reason error)
)

// UpgradeRequests will upgrade any incoming websocket request
// matching the given pattern, and then call the event handler
// function with connection events.
func UpgradeRequests(pattern string, eventHandler EventHandler) {
	upgrader := websocket.Upgrader{
		HandshakeTimeout: HandshakeTimeout,
		ReadBufferSize:   ReadBufferSize,
		WriteBufferSize:  WriteBufferSize,
		Subprotocols:     []string{"birect"},
		Error:            ErrorFn,
		CheckOrigin:      CheckOrigin,
	}
	http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
		wsConn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			_checkAndGenerateEvent(eventHandler, Error, nil, nil, err)
		} else if wsConn.Subprotocol() != "birect" {
			_checkAndGenerateEvent(eventHandler, Error, nil, nil, errs.New(nil, "Unable to negotiate birect subprotocol"))
		} else {
			newConn(r, wsConn, eventHandler)
		}
	})
	return
}
<commit_msg>Ensure that incoming request version is 13 before proceeding<commit_after>package ws

import (
	"net/http"

	"github.com/gorilla/websocket"
	"github.com/marcuswestin/go-errs"
)

var (
	/*
		HTTP Upgrade Options

		These are the different options
	*/

	// CheckOrigin will be called by the upgrader on incoming requests to check
	// the origin. If the CheckOrigin function returns false, then the Upgrade
	// method fails the WebSocket handshake with HTTP status 403.
	// If nil, then use a safe default: fail the handshake if the Origin request
	// header is present and not equal to the Host request header.
	CheckOrigin func(r *http.Request) bool

	// ErrorFn specifies the function for generating HTTP error responses.
	// If nil, then http.Error is used to generate the HTTP response.
	ErrorFn func(w http.ResponseWriter, r *http.Request, status int, reason error)
)

// UpgradeRequests will upgrade any incoming websocket request
// matching the given pattern, and then call the event handler
// function with connection events.
func UpgradeRequests(pattern string, eventHandler EventHandler) {
	upgrader := websocket.Upgrader{
		HandshakeTimeout: HandshakeTimeout,
		ReadBufferSize:   ReadBufferSize,
		WriteBufferSize:  WriteBufferSize,
		Subprotocols:     []string{"birect"},
		Error:            ErrorFn,
		CheckOrigin:      CheckOrigin,
	}
	http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
		if version := r.Header.Get("Sec-Websocket-Version"); version != "13" {
			http.Error(w, "Sec-Websocket-Version must be 13", 400)
			return
		}
		wsConn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			_checkAndGenerateEvent(eventHandler, Error, nil, nil, err)
		} else if wsConn.Subprotocol() != "birect" {
			_checkAndGenerateEvent(eventHandler, Error, nil, nil, errs.New(nil, "Unable to negotiate birect subprotocol"))
		} else {
			newConn(r, wsConn, eventHandler)
		}
	})
	return
}
<|endoftext|>


<commit_before>package address

import (
	"github.com/eyecuelab/kit/str"
	"github.com/eyecuelab/kit/stringslice"
)

//TotalDistance returns the sum of the levenshtein distance of the components of the addresses.
func TotalDistance(placeA, placeB *Address) (distance int) {
	a, b := placeA.StringSlice(), placeB.StringSlice()
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//SharedComponentDistance returns the sum of the levenshtein distance of the shared components of the addresses
func SharedComponentDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSliceFromNonempty(), filteredB.StringSliceFromNonempty()
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//NormalizedTotalDistance Performs heavy normalization on both addresses, then compares the distance.
func NormalizedTotalDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSlice(), filteredB.StringSlice()
	a, b = str.Map(normalize, a), str.Map(normalize, b)
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//NormalizedSharedComponentDistance performs heavy normalization on both addresses, then compares the distance.
func NormalizedSharedComponentDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSliceFromNonempty(), filteredB.StringSliceFromNonempty()
	a, b = str.Map(normalize, a), str.Map(normalize, b)
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}
<commit_msg>changed name for exported func<commit_after>package address

import (
	"github.com/eyecuelab/kit/str"
	"github.com/eyecuelab/kit/stringslice"
)

//TotalDistance returns the sum of the levenshtein distance of the components of the addresses.
func TotalDistance(placeA, placeB *Address) (distance int) {
	a, b := placeA.StringSlice(), placeB.StringSlice()
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//SharedComponentDistance returns the sum of the levenshtein distance of the shared components of the addresses
func SharedComponentDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSliceFromNonempty(), filteredB.StringSliceFromNonempty()
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//NormalizedTotalDistance Performs heavy normalization on both addresses, then compares the distance.
func NormalizedTotalDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSlice(), filteredB.StringSlice()
	a, b = str.Map(Normalize, a), str.Map(Normalize, b)
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}

//NormalizedSharedComponentDistance performs heavy normalization on both addresses, then compares the distance.
func NormalizedSharedComponentDistance(placeA, placeB *Address) (distance int) {
	filteredA, filteredB := SharedComponents(*placeA, *placeB)
	a, b := filteredA.StringSliceFromNonempty(), filteredB.StringSliceFromNonempty()
	a, b = str.Map(Normalize, a), str.Map(Normalize, b)
	distance, _ = stringslice.TotalDistance(a, b)
	return distance
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	ggio "github.com/gogo/protobuf/io"
	p2p_crypto "github.com/ipfs/go-libp2p-crypto"
	p2p_peer "github.com/ipfs/go-libp2p-peer"
	p2p_pstore "github.com/ipfs/go-libp2p-peerstore"
	multiaddr "github.com/jbenet/go-multiaddr"
	p2p_host "github.com/libp2p/go-libp2p/p2p/host"
	p2p_net "github.com/libp2p/go-libp2p/p2p/net"
	mc "github.com/mediachain/concat/mc"
	pb "github.com/mediachain/concat/proto"
	"log"
	"os"
	"time"
)

type Node struct {
	id    p2p_peer.ID
	privk p2p_crypto.PrivKey
	host  p2p_host.Host
}

func (node *Node) pingHandler(s p2p_net.Stream) {

}

func (node *Node) registerPeer(dir p2p_pstore.PeerInfo, addrs ...multiaddr.Multiaddr) {
	// directory failure is a fatality for now
	ctx := context.Background()

	err := node.host.Connect(ctx, dir)
	if err != nil {
		log.Printf("Failed to connect to directory")
		log.Fatal(err)
	}

	s, err := node.host.NewStream(ctx, dir.ID, "/mediachain/dir/register")
	if err != nil {
		log.Printf("Failed to open directory stream")
		log.Fatal(err)
	}
	defer s.Close()

	pbpi := pb.PeerInfo{}
	pbpi.Id = string(node.id)
	pbpi.Addr = make([][]byte, len(addrs))
	for x, addr := range addrs {
		pbpi.Addr[x] = addr.Bytes()
	}
	msg := pb.RegisterPeer{&pbpi}

	w := ggio.NewDelimitedWriter(s)
	for {
		log.Printf("registering with directory")
		err = w.WriteMsg(&msg)
		if err != nil {
			log.Printf("Failed to register with directory")
			log.Fatal(err)
		}

		time.Sleep(5 * time.Minute)
	}
}

func main() {
	port := flag.Int("l", 9001, "Listen port")
	flag.Parse()

	if len(flag.Args()) != 1 {
		fmt.Printf("Usage: %s [-l port] directory\n", os.Args[0])
		os.Exit(1)
	}

	addr, err := mc.ParseAddress(fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", *port))
	if err != nil {
		log.Fatal(err)
	}

	dir, err := mc.ParseHandle(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Generating key pair\n")
	privk, pubk, err := mc.GenerateKeyPair()
	if err != nil {
		log.Fatal(err)
	}

	id, err := p2p_peer.IDFromPublicKey(pubk)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("ID: %s\n", id.Pretty())

	host, err := mc.NewHost(privk, addr)
	if err != nil {
		log.Fatal(err)
	}

	node := &Node{id: id, privk: privk, host: host}
	host.SetStreamHandler("/mediachain/node/ping", node.pingHandler)
	go node.registerPeer(dir, addr)

	log.Printf("I am %s/%s", addr, id.Pretty())
	select {}
}
<commit_msg>mcnode: add dir to node structure<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	ggio "github.com/gogo/protobuf/io"
	p2p_crypto "github.com/ipfs/go-libp2p-crypto"
	p2p_peer "github.com/ipfs/go-libp2p-peer"
	p2p_pstore "github.com/ipfs/go-libp2p-peerstore"
	multiaddr "github.com/jbenet/go-multiaddr"
	p2p_host "github.com/libp2p/go-libp2p/p2p/host"
	p2p_net "github.com/libp2p/go-libp2p/p2p/net"
	mc "github.com/mediachain/concat/mc"
	pb "github.com/mediachain/concat/proto"
	"log"
	"os"
	"time"
)

type Node struct {
	id    p2p_peer.ID
	privk p2p_crypto.PrivKey
	host  p2p_host.Host
	dir   p2p_pstore.PeerInfo
}

func (node *Node) pingHandler(s p2p_net.Stream) {

}

func (node *Node) registerPeer(addrs ...multiaddr.Multiaddr) {
	// directory failure is a fatality for now
	ctx := context.Background()

	err := node.host.Connect(ctx, node.dir)
	if err != nil {
		log.Printf("Failed to connect to directory")
		log.Fatal(err)
	}

	s, err := node.host.NewStream(ctx, node.dir.ID, "/mediachain/dir/register")
	if err != nil {
		log.Printf("Failed to open directory stream")
		log.Fatal(err)
	}
	defer s.Close()

	pbpi := pb.PeerInfo{}
	pbpi.Id = string(node.id)
	pbpi.Addr = make([][]byte, len(addrs))
	for x, addr := range addrs {
		pbpi.Addr[x] = addr.Bytes()
	}
	msg := pb.RegisterPeer{&pbpi}

	w := ggio.NewDelimitedWriter(s)
	for {
		log.Printf("registering with directory")
		err = w.WriteMsg(&msg)
		if err != nil {
			log.Printf("Failed to register with directory")
			log.Fatal(err)
		}

		time.Sleep(5 * time.Minute)
	}
}

func main() {
	port := flag.Int("l", 9001, "Listen port")
	flag.Parse()

	if len(flag.Args()) != 1 {
		fmt.Printf("Usage: %s [-l port] directory\n", os.Args[0])
		os.Exit(1)
	}

	addr, err := mc.ParseAddress(fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", *port))
	if err != nil {
		log.Fatal(err)
	}

	dir, err := mc.ParseHandle(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Generating key pair\n")
	privk, pubk, err := mc.GenerateKeyPair()
	if err != nil {
		log.Fatal(err)
	}

	id, err := p2p_peer.IDFromPublicKey(pubk)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("ID: %s\n", id.Pretty())

	host, err := mc.NewHost(privk, addr)
	if err != nil {
		log.Fatal(err)
	}

	node := &Node{id: id, privk: privk, host: host, dir: dir}
	host.SetStreamHandler("/mediachain/node/ping", node.pingHandler)
	go node.registerPeer(addr)

	log.Printf("I am %s/%s", addr, id.Pretty())
	select {}
}
<|endoftext|>


<commit_before>/*
 * This file is part of the libvirt-go-xml project
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Copyright (C) 2016 Red Hat, Inc.
 *
 */

package libvirtxml

import (
	"encoding/xml"
)

type DomainController struct {
	Type  string `xml:"type,attr"`
	Index string `xml:"index,attr"`
}

type DomainDiskFileSource struct {
	File string `xml:"file,attr"`
}

type DomainDiskDriver struct {
	Name string `xml:"name,attr"`
	Type string `xml:"type,attr"`
}

type DomainDiskTarget struct {
	Dev string `xml:"dev,attr"`
	Bus string `xml:"bus,attr"`
}

type DomainDisk struct {
	Type       string                `xml:"type,attr"`
	Device     string                `xml:"device,attr"`
	Snapshot   string                `xml:"snapshot,attr,omitempty"`
	Driver     DomainDiskDriver      `xml:"driver"`
	FileSource *DomainDiskFileSource `xml:"source"`
	Target     *DomainDiskTarget     `xml:"target"`
}

type DomainInterfaceMAC struct {
	Address string `xml:"address,attr"`
}

type DomainInterfaceModel struct {
	Type string `xml:type,attr"`
}

type DomainInterface struct {
	Type  string               `xml:"type,attr"`
	MAC   string               `xml:"mac"`
	Model DomainInterfaceModel `xml:"model"`
}

type DomainChardev struct {
	Type string `xml:"type,attr"`
}

type DomainInput struct {
	Type string `xml:"type,attr"`
	Bus  string `xml:"bus,attr"`
}

type DomainGraphicListener struct {
	Type    string `xml:"type,attr"`
	Address string `xml:"address,attr,omitempty"`
	Network string `xml:"network,attr,omitempty"`
	Socket  string `xml:"socket,attr,omitempty"`
}

type DomainGraphic struct {
	Type          string                  `xml:"type,attr"`
	AutoPort      string                  `xml:"autoport,attr,omitempty"`
	Port          int                     `xml:"port,attr,omitempty"`
	TLSPort       int                     `xml:"tlsPort,attr,omitempty"`
	WebSocket     int                     `xml:"websocket,attr,omitempty"`
	Listen        string                  `xml:"listen,attr,omitempty"`
	Socket        string                  `xml:"socket,attr,omitempty"`
	Keymap        string                  `xml:"keymap,attr,omitempty"`
	Passwd        string                  `xml:"passwd,attr,omitempty"`
	PasswdValidTo string                  `xml:"passwdValidTo,attr,omitempty"`
	Connected     string                  `xml:"connected,attr,omitempty"`
	SharePolicy   string                  `xml:"sharePolicy,attr,omitempty"`
	DefaultMode   string                  `xml:"defaultMode,attr,omitempty"`
	Display       string                  `xml:"display,attr,omitempty"`
	XAuth         string                  `xml:"xauth,attr,omitempty"`
	FullScreen    string                  `xml:"fullscreen,attr,omitempty"`
	ReplaceUser   string                  `xml:"replaceUser,attr,omitempty"`
	MultiUser     string                  `xml:"multiUser,attr,omitempty"`
	Listeners     []DomainGraphicListener `xml:"listen"`
}

type DomainVideoModel struct {
	Type string `xml:"type,attr"`
}

type DomainVideo struct {
	Model DomainVideoModel `xml:"model"`
}

type DomainDeviceList struct {
	Controllers []DomainController `xml:"controller"`
	Disks       []DomainDisk       `xml:"disk"`
	Interfaces  []DomainInterface  `xml:"interface"`
	Serials     []DomainChardev    `xml:"serial"`
	Consoles    []DomainChardev    `xml:"console"`
	Inputs      []DomainInput      `xml:"input"`
	Graphics    []DomainGraphic    `xml:"graphics"`
	Videos      []DomainVideo      `xml:"video"`
}

type DomainMemory struct {
	Value int    `xml:",chardata"`
	Unit  string `xml:"unit,attr,omitempty"`
}

type DomainMaxMemory struct {
	Value int    `xml:",chardata"`
	Unit  string `xml:"unit,attr,omitempty"`
	Slots int    `xml:"slots,attr,omitempty"`
}

type DomainOSType struct {
	Arch    string `xml:"arch,attr"`
	Machine string `xml:"machine,attr"`
	Type    string `xml:",chardata"`
}

type DomainSMBios struct {
	Mode string `xml:"mode,attr"`
}

type DomainNVRam struct {
	NVRam    string `xml:",chardata"`
	Template string `xml:"template,attr"`
}

type DomainBootDevice struct {
	Dev string `xml:"dev,attr"`
}

type DomainBootMenu struct {
	Enabled string `xml:"enabled,attr"`
	Timeout string `xml:"timeout,attr"`
}

type DomainSysInfo struct {
	Type      string               `xml:"type,attr"`
	System    []DomainSysInfoEntry `xml:"system>entry"`
	BIOS      []DomainSysInfoEntry `xml:"bios>entry"`
	BaseBoard []DomainSysInfoEntry `xml:"baseBoard>entry"`
}

type DomainSysInfoEntry struct {
	Name  string `xml:"name,attr"`
	Value string `xml:",chardata"`
}

type DomainBIOS struct {
	UseSerial     string `xml:"useserial,attr"`
	RebootTimeout string `xml:"rebootTimeout,attr"`
}

type DomainLoader struct {
	Path     string `xml:",chardata"`
	Readonly string `xml:"readonly,attr"`
	Secure   string `xml:"secure,attr"`
	Type     string `xml:"type,attr"`
}

type DomainOS struct {
	Type        *DomainOSType      `xml:"type"`
	Loader      *DomainLoader      `xml:"loader"`
	NVRam       *DomainNVRam       `xml:"nvram"`
	Kernel      string             `xml:"kernel,omitempty"`
	Initrd      string             `xml:"initrd,omitempty"`
	KernelArgs  string             `xml:"cmdline,omitempty"`
	BootDevices []DomainBootDevice `xml:"boot"`
	BootMenu    *DomainBootMenu    `xml:"bootmenu"`
	SMBios      *DomainSMBios      `xml:"smbios"`
	BIOS        *DomainBIOS        `xml:"bios"`
	Init        string             `xml:"init,omitempty"`
	InitArgs    []string           `xml:"initarg"`
}

type Domain struct {
	XMLName       xml.Name          `xml:"domain"`
	Type          string            `xml:"type,attr,omitempty"`
	Name          string            `xml:"name"`
	UUID          *string           `xml:"uuid"`
	Memory        *DomainMemory     `xml:"memory"`
	CurrentMemory *DomainMemory     `xml:"currentMemory"`
	MaximumMemory *DomainMaxMemory  `xml:"maxMemory"`
	Devices       *DomainDeviceList `xml:"devices"`
	OS            *DomainOS         `xml:"os"`
	SysInfo       *DomainSysInfo    `xml:"sysinfo"`
	OnPoweroff    string            `xml:"on_poweroff,omitempty"`
	OnReboot      string            `xml:"on_reboot,omitempty"`
	OnCrash       string            `xml:"on_crash,omitempty"`
}

func (d *Domain) Unmarshal(doc string) error {
	return xml.Unmarshal([]byte(doc), d)
}

func (d *Domain) Marshal() (string, error) {
	doc, err := xml.MarshalIndent(d, "", "  ")
	if err != nil {
		return "", err
	}
	return string(doc), nil
}
<commit_msg>Make disk driver source optional<commit_after>/*
 * This file is part of the libvirt-go-xml project
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Copyright (C) 2016 Red Hat, Inc.
 *
 */

package libvirtxml

import (
	"encoding/xml"
)

type DomainController struct {
	Type  string `xml:"type,attr"`
	Index string `xml:"index,attr"`
}

type DomainDiskFileSource struct {
	File string `xml:"file,attr"`
}

type DomainDiskDriver struct {
	Name string `xml:"name,attr"`
	Type string `xml:"type,attr"`
}

type DomainDiskTarget struct {
	Dev string `xml:"dev,attr"`
	Bus string `xml:"bus,attr"`
}

type DomainDisk struct {
	Type       string                `xml:"type,attr"`
	Device     string                `xml:"device,attr"`
	Snapshot   string                `xml:"snapshot,attr,omitempty"`
	Driver     *DomainDiskDriver     `xml:"driver"`
	FileSource *DomainDiskFileSource `xml:"source"`
	Target     *DomainDiskTarget     `xml:"target"`
}

type DomainInterfaceMAC struct {
	Address string `xml:"address,attr"`
}

type DomainInterfaceModel struct {
	Type string `xml:type,attr"`
}

type DomainInterface struct {
	Type  string               `xml:"type,attr"`
	MAC   string               `xml:"mac"`
	Model DomainInterfaceModel `xml:"model"`
}

type DomainChardev struct {
	Type string `xml:"type,attr"`
}

type DomainInput struct {
	Type string `xml:"type,attr"`
	Bus  string `xml:"bus,attr"`
}

type DomainGraphicListener struct {
	Type    string `xml:"type,attr"`
	Address string `xml:"address,attr,omitempty"`
	Network string `xml:"network,attr,omitempty"`
	Socket  string `xml:"socket,attr,omitempty"`
}

type DomainGraphic struct {
	Type          string                  `xml:"type,attr"`
	AutoPort      string                  `xml:"autoport,attr,omitempty"`
	Port          int                     `xml:"port,attr,omitempty"`
	TLSPort       int                     `xml:"tlsPort,attr,omitempty"`
	WebSocket     int                     `xml:"websocket,attr,omitempty"`
	Listen        string                  `xml:"listen,attr,omitempty"`
	Socket        string                  `xml:"socket,attr,omitempty"`
	Keymap        string                  `xml:"keymap,attr,omitempty"`
	Passwd        string                  `xml:"passwd,attr,omitempty"`
	PasswdValidTo string                  `xml:"passwdValidTo,attr,omitempty"`
	Connected     string                  `xml:"connected,attr,omitempty"`
	SharePolicy   string                  `xml:"sharePolicy,attr,omitempty"`
	DefaultMode   string                  `xml:"defaultMode,attr,omitempty"`
	Display       string                  `xml:"display,attr,omitempty"`
	XAuth         string                  `xml:"xauth,attr,omitempty"`
	FullScreen    string                  `xml:"fullscreen,attr,omitempty"`
	ReplaceUser   string                  `xml:"replaceUser,attr,omitempty"`
	MultiUser     string                  `xml:"multiUser,attr,omitempty"`
	Listeners     []DomainGraphicListener `xml:"listen"`
}

type DomainVideoModel struct {
	Type string `xml:"type,attr"`
}

type DomainVideo struct {
	Model DomainVideoModel `xml:"model"`
}

type DomainDeviceList struct {
	Controllers []DomainController `xml:"controller"`
	Disks       []DomainDisk       `xml:"disk"`
	Interfaces  []DomainInterface  `xml:"interface"`
	Serials     []DomainChardev    `xml:"serial"`
	Consoles    []DomainChardev    `xml:"console"`
	Inputs      []DomainInput      `xml:"input"`
	Graphics    []DomainGraphic    `xml:"graphics"`
	Videos      []DomainVideo      `xml:"video"`
}

type DomainMemory struct {
	Value int    `xml:",chardata"`
	Unit  string `xml:"unit,attr,omitempty"`
}

type DomainMaxMemory struct {
	Value int    `xml:",chardata"`
	Unit  string `xml:"unit,attr,omitempty"`
	Slots int    `xml:"slots,attr,omitempty"`
}

type DomainOSType struct {
	Arch    string `xml:"arch,attr"`
	Machine string `xml:"machine,attr"`
	Type    string `xml:",chardata"`
}

type DomainSMBios struct {
	Mode string `xml:"mode,attr"`
}

type DomainNVRam struct {
	NVRam    string `xml:",chardata"`
	Template string `xml:"template,attr"`
}

type DomainBootDevice struct {
	Dev string `xml:"dev,attr"`
}

type DomainBootMenu struct {
	Enabled string `xml:"enabled,attr"`
	Timeout string `xml:"timeout,attr"`
}

type DomainSysInfo struct {
	Type      string               `xml:"type,attr"`
	System    []DomainSysInfoEntry `xml:"system>entry"`
	BIOS      []DomainSysInfoEntry `xml:"bios>entry"`
	BaseBoard []DomainSysInfoEntry `xml:"baseBoard>entry"`
}

type DomainSysInfoEntry struct {
	Name  string `xml:"name,attr"`
	Value string `xml:",chardata"`
}

type DomainBIOS struct {
	UseSerial     string `xml:"useserial,attr"`
	RebootTimeout string `xml:"rebootTimeout,attr"`
}

type DomainLoader struct {
	Path     string `xml:",chardata"`
	Readonly string `xml:"readonly,attr"`
	Secure   string `xml:"secure,attr"`
	Type     string `xml:"type,attr"`
}

type DomainOS struct {
	Type        *DomainOSType      `xml:"type"`
	Loader      *DomainLoader      `xml:"loader"`
	NVRam       *DomainNVRam       `xml:"nvram"`
	Kernel      string             `xml:"kernel,omitempty"`
	Initrd      string             `xml:"initrd,omitempty"`
	KernelArgs  string             `xml:"cmdline,omitempty"`
	BootDevices []DomainBootDevice `xml:"boot"`
	BootMenu    *DomainBootMenu    `xml:"bootmenu"`
	SMBios      *DomainSMBios      `xml:"smbios"`
	BIOS        *DomainBIOS        `xml:"bios"`
	Init        string             `xml:"init,omitempty"`
	InitArgs    []string           `xml:"initarg"`
}

type Domain struct {
	XMLName       xml.Name          `xml:"domain"`
	Type          string            `xml:"type,attr,omitempty"`
	Name          string            `xml:"name"`
	UUID          *string           `xml:"uuid"`
	Memory        *DomainMemory     `xml:"memory"`
	CurrentMemory *DomainMemory     `xml:"currentMemory"`
	MaximumMemory *DomainMaxMemory  `xml:"maxMemory"`
	Devices       *DomainDeviceList `xml:"devices"`
	OS            *DomainOS         `xml:"os"`
	SysInfo       *DomainSysInfo    `xml:"sysinfo"`
	OnPoweroff    string            `xml:"on_poweroff,omitempty"`
	OnReboot      string            `xml:"on_reboot,omitempty"`
	OnCrash       string            `xml:"on_crash,omitempty"`
}

func (d *Domain) Unmarshal(doc string) error {
	return xml.Unmarshal([]byte(doc), d)
}

func (d *Domain) Marshal() (string, error) {
	doc, err := xml.MarshalIndent(d, "", "  ")
	if err != nil {
		return "", err
	}
	return string(doc), nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package json

import (
	"github.com/stretchr/testify/assert"
	"os"
	"testing"
)

func TestNewWithInaccessibleFilename(t *testing.T) {
	p := string("/this_should_not_exist/badplacetolog.log")
	collector, err := New(p)
	assert.NotEqual(t, err, error(nil))
	assert.Equal(t, collector, (*Collector)(nil))
}

func TestNewWithFileName(t *testing.T) {
	p := string("/tmp/okplacetolog.log")
	defer func() { _ = os.Remove(p) }()
	collector, err := New(p)
	assert.Equal(t, err, error(nil))
	assert.NotEqual(t, collector, (*Collector)(nil))
}

func TestNewWithLocalFilename1(t *testing.T) {
	p := string("./okplacetolog.log")
	defer func() { _ = os.Remove(p) }()
	collector, err := New(p)
	assert.Equal(t, err, error(nil))
	assert.NotEqual(t, collector, (*Collector)(nil))
}

func TestNewWithLocalFilename2(t *testing.T) {
	p := string("okplacetolog.log")
	defer func() { _ = os.Remove(p) }()
	collector, err := New(p)
	assert.Equal(t, err, error(nil))
	assert.NotEqual(t, collector, (*Collector)(nil))
}
<commit_msg>[test/refactor] Use t.Run() to reduce boilerplate<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package json

import (
	"github.com/stretchr/testify/assert"
	"os"
	"testing"
)

func TestNew(t *testing.T) {
	testdata := map[string]bool{
		"/nonexistent/badplacetolog.log":   false,
		os.TempDir() + "/okplacetolog.log": true,
		"./okplacetolog.log":               true,
		"okplacetolog.log":                 true,
	}

	for path, succ := range testdata {
		t.Run("path="+path, func(t *testing.T) {
			defer func() { _ = os.Remove(path) }()

			collector, err := New(path)
			if succ {
				assert.NoError(t, err)
				assert.NotNil(t, collector)
			} else {
				assert.Error(t, err)
				assert.Nil(t, collector)
			}
		})
	}
}
<|endoftext|>


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

package sqlite

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"log"
	"os"
	"reflect"
	"time"
	"unsafe"
)

func init() {
	sql.Register("sqlite3", &Driver{})
	if os.Getenv("SQLITE_LOG") != "" {
		ConfigLog(func(d interface{}, err error, msg string) {
			log.Printf("%s: %s, %s\n", d, err, msg)
		}, "SQLITE")
	}
	ConfigMemStatus(false)
}

// Adapter to database/sql/driver
type Driver struct {
}
type connImpl struct {
	c *Conn
}
type stmtImpl struct {
	s            *Stmt
	rowsRef      bool // true if there is a rowsImpl associated to this statement that has not been closed.
	pendingClose bool
}
type rowsImpl struct {
	s           *stmtImpl
	columnNames []string // cache
}

// Open opens a new database connection.
// ":memory:" for memory db,
// "" for temp file db
func (d *Driver) Open(name string) (driver.Conn, error) {
	c, err := Open(name, OpenNoMutex, OpenReadWrite, OpenCreate)
	if err != nil {
		return nil, err
	}
	c.BusyTimeout(time.Duration(10) * time.Second)
	return &connImpl{c}, nil
}

// PRAGMA schema_version may be used to detect when the database schema is altered

func (c *connImpl) Exec(query string, args []driver.Value) (driver.Result, error) {
	// https://code.google.com/p/go-wiki/wiki/cgo#Turning_C_arrays_into_Go_slices
	var iargs []interface{}
	if len(args) > 0 {
		h := (*reflect.SliceHeader)(unsafe.Pointer(&iargs))
		h.Data = uintptr(unsafe.Pointer(&args[0]))
		h.Len = len(args)
		h.Cap = cap(args)
	}
	if err := c.c.Exec(query, iargs...); err != nil {
		return nil, err
	}
	return c, nil // FIXME RowAffected/noRows
}

// TODO How to know that the last Stmt has done an INSERT? An authorizer?
func (c *connImpl) LastInsertId() (int64, error) {
	return c.c.LastInsertRowid(), nil
}

// TODO How to know that the last Stmt has done a DELETE/INSERT/UPDATE? An authorizer?
func (c *connImpl) RowsAffected() (int64, error) {
	return int64(c.c.Changes()), nil
}

func (c *connImpl) Prepare(query string) (driver.Stmt, error) {
	s, err := c.c.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &stmtImpl{s: s}, nil
}

func (c *connImpl) Close() error {
	return c.c.Close()
}

func (c *connImpl) Begin() (driver.Tx, error) {
	if err := c.c.Begin(); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *connImpl) Commit() error {
	return c.c.Commit()
}
func (c *connImpl) Rollback() error {
	return c.c.Rollback()
}

func (s *stmtImpl) Close() error {
	if s.rowsRef { // Currently, it never happens because the sql.Stmt doesn't call driver.Stmt in this case
		s.pendingClose = true
		return nil
	}
	return s.s.Finalize()
}

func (s *stmtImpl) NumInput() int {
	return s.s.BindParameterCount()
}

func (s *stmtImpl) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	if err := s.s.exec(); err != nil {
		return nil, err
	}
	return s, nil // FIXME RowAffected/noRows
}

// TODO How to know that this Stmt has done an INSERT? An authorizer?
func (s *stmtImpl) LastInsertId() (int64, error) {
	return s.s.c.LastInsertRowid(), nil
}

// TODO How to know that this Stmt has done a DELETE/INSERT/UPDATE? An authorizer?
func (s *stmtImpl) RowsAffected() (int64, error) {
	return int64(s.s.c.Changes()), nil
}

func (s *stmtImpl) Query(args []driver.Value) (driver.Rows, error) {
	if s.rowsRef {
		return nil, errors.New("Previously returned Rows still not closed")
	}
	if err := s.bind(args); err != nil {
		return nil, err
	}
	s.rowsRef = true
	return &rowsImpl{s, nil}, nil
}

func (s *stmtImpl) bind(args []driver.Value) error {
	for i, v := range args {
		if err := s.s.BindByIndex(i+1, v); err != nil {
			return err
		}
	}
	return nil
}

func (r *rowsImpl) Columns() []string {
	if r.columnNames == nil {
		r.columnNames = r.s.s.ColumnNames()
	}
	return r.columnNames
}

func (r *rowsImpl) Next(dest []driver.Value) error {
	ok, err := r.s.s.Next()
	if err != nil {
		return err
	}
	if !ok {
		return io.EOF
	}
	for i := range dest {
		value := r.s.s.ScanValue(i)
		switch value := value.(type) {
		case string: // "All string values must be converted to []byte."
			dest[i] = []byte(value)
		default:
			dest[i] = value
		}
	}
	return nil
}

func (r *rowsImpl) Close() error {
	r.s.rowsRef = false
	if r.s.pendingClose {
		return r.s.Close()
	}
	return r.s.s.Reset()
}
<commit_msg>Adds support to URI in driver.Open.<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlite

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"log"
	"os"
	"reflect"
	"time"
	"unsafe"
)

func init() {
	sql.Register("sqlite3", &Driver{})
	if os.Getenv("SQLITE_LOG") != "" {
		ConfigLog(func(d interface{}, err error, msg string) {
			log.Printf("%s: %s, %s\n", d, err, msg)
		}, "SQLITE")
	}
	ConfigMemStatus(false)
}

// Adapter to database/sql/driver
type Driver struct {
}
type connImpl struct {
	c *Conn
}
type stmtImpl struct {
	s            *Stmt
	rowsRef      bool // true if there is a rowsImpl associated to this statement that has not been closed.
	pendingClose bool
}
type rowsImpl struct {
	s           *stmtImpl
	columnNames []string // cache
}

// Open opens a new database connection.
// ":memory:" for memory db,
// "" for temp file db
func (d *Driver) Open(name string) (driver.Conn, error) {
	c, err := Open(name, OpenUri, OpenNoMutex, OpenReadWrite, OpenCreate)
	if err != nil {
		return nil, err
	}
	c.BusyTimeout(time.Duration(10) * time.Second)
	return &connImpl{c}, nil
}

// PRAGMA schema_version may be used to detect when the database schema is altered

func (c *connImpl) Exec(query string, args []driver.Value) (driver.Result, error) {
	// https://code.google.com/p/go-wiki/wiki/cgo#Turning_C_arrays_into_Go_slices
	var iargs []interface{}
	if len(args) > 0 {
		h := (*reflect.SliceHeader)(unsafe.Pointer(&iargs))
		h.Data = uintptr(unsafe.Pointer(&args[0]))
		h.Len = len(args)
		h.Cap = cap(args)
	}
	if err := c.c.Exec(query, iargs...); err != nil {
		return nil, err
	}
	return c, nil // FIXME RowAffected/noRows
}

// TODO How to know that the last Stmt has done an INSERT? An authorizer?
func (c *connImpl) LastInsertId() (int64, error) {
	return c.c.LastInsertRowid(), nil
}

// TODO How to know that the last Stmt has done a DELETE/INSERT/UPDATE? An authorizer?
func (c *connImpl) RowsAffected() (int64, error) {
	return int64(c.c.Changes()), nil
}

func (c *connImpl) Prepare(query string) (driver.Stmt, error) {
	s, err := c.c.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &stmtImpl{s: s}, nil
}

func (c *connImpl) Close() error {
	return c.c.Close()
}

func (c *connImpl) Begin() (driver.Tx, error) {
	if err := c.c.Begin(); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *connImpl) Commit() error {
	return c.c.Commit()
}
func (c *connImpl) Rollback() error {
	return c.c.Rollback()
}

func (s *stmtImpl) Close() error {
	if s.rowsRef { // Currently, it never happens because the sql.Stmt doesn't call driver.Stmt in this case
		s.pendingClose = true
		return nil
	}
	return s.s.Finalize()
}

func (s *stmtImpl) NumInput() int {
	return s.s.BindParameterCount()
}

func (s *stmtImpl) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	if err := s.s.exec(); err != nil {
		return nil, err
	}
	return s, nil // FIXME RowAffected/noRows
}

// TODO How to know that this Stmt has done an INSERT? An authorizer?
func (s *stmtImpl) LastInsertId() (int64, error) {
	return s.s.c.LastInsertRowid(), nil
}

// TODO How to know that this Stmt has done a DELETE/INSERT/UPDATE? An authorizer?
func (s *stmtImpl) RowsAffected() (int64, error) {
	return int64(s.s.c.Changes()), nil
}

func (s *stmtImpl) Query(args []driver.Value) (driver.Rows, error) {
	if s.rowsRef {
		return nil, errors.New("Previously returned Rows still not closed")
	}
	if err := s.bind(args); err != nil {
		return nil, err
	}
	s.rowsRef = true
	return &rowsImpl{s, nil}, nil
}

func (s *stmtImpl) bind(args []driver.Value) error {
	for i, v := range args {
		if err := s.s.BindByIndex(i+1, v); err != nil {
			return err
		}
	}
	return nil
}

func (r *rowsImpl) Columns() []string {
	if r.columnNames == nil {
		r.columnNames = r.s.s.ColumnNames()
	}
	return r.columnNames
}

func (r *rowsImpl) Next(dest []driver.Value) error {
	ok, err := r.s.s.Next()
	if err != nil {
		return err
	}
	if !ok {
		return io.EOF
	}
	for i := range dest {
		value := r.s.s.ScanValue(i)
		switch value := value.(type) {
		case string: // "All string values must be converted to []byte."
			dest[i] = []byte(value)
		default:
			dest[i] = value
		}
	}
	return nil
}

func (r *rowsImpl) Close() error {
	r.s.rowsRef = false
	if r.s.pendingClose {
		return r.s.Close()
	}
	return r.s.s.Reset()
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// The driver should be used via the database/sql package:
//
//  import "database/sql"
//  import _ "github.com/go-sql-driver/mysql"
//
//  db, err := sql.Open("mysql", "user:password@/dbname")
//
// See https://github.com/go-sql-driver/mysql#usage for details
package mysql

import (
	"database/sql"
	"database/sql/driver"
	"net"
)

// This struct is exported to make the driver directly accessible.
// In general the driver is used via the database/sql package.
type MySQLDriver struct{}

// DialFunc is a function which can be used to establish the network connection.
// Custom dial functions must be registered with RegisterDial
type DialFunc func(addr string) (net.Conn, error)

var dials map[string]DialFunc

// RegisterDial registers a custom dial function. It can then be used by the
// network address mynet(addr), where mynet is the registered new network.
// addr is passed as a parameter to the dial function.
func RegisterDial(net string, dial DialFunc) {
	if dials == nil {
		dials = make(map[string]DialFunc)
	}
	dials[net] = dial
}

// Open new Connection.
// See https://github.com/go-sql-driver/mysql#dsn-data-source-name for how
// the DSN string is formated
func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
	var err error

	// New mysqlConn
	mc := &mysqlConn{
		maxPacketAllowed: maxPacketSize,
		maxWriteSize:     maxPacketSize - 1,
	}
	mc.cfg, err = parseDSN(dsn)
	if err != nil {
		return nil, err
	}

	// Connect to Server
	if dial, ok := dials[mc.cfg.net]; ok {
		mc.netConn, err = dial(mc.cfg.addr)
	} else {
		nd := net.Dialer{Timeout: mc.cfg.timeout}
		mc.netConn, err = nd.Dial(mc.cfg.net, mc.cfg.addr)
	}
	if err != nil {
		return nil, err
	}

	// Enable TCP Keepalives on TCP connections
	if tc, ok := mc.netConn.(*net.TCPConn); ok {
		if err := tc.SetKeepAlive(true); err != nil {
			// Don't send COM_QUIT before handshake.
			mc.netConn.Close()
			mc.netConn = nil
			return nil, err
		}
	}

	mc.buf = newBuffer(mc.netConn)

	// Reading Handshake Initialization Packet
	cipher, err := mc.readInitPacket()
	if err != nil {
		mc.Close()
		return nil, err
	}

	// Send Client Authentication Packet
	if err = mc.writeAuthPacket(cipher); err != nil {
		mc.Close()
		return nil, err
	}

	// Read Result Packet
	err = mc.readResultOK()
	if err != nil {
		// Retry with old authentication method, if allowed
		if mc.cfg != nil && mc.cfg.allowOldPasswords && err == ErrOldPassword {
			if err = mc.writeOldAuthPacket(cipher); err != nil {
				mc.Close()
				return nil, err
			}
			if err = mc.readResultOK(); err != nil {
				mc.Close()
				return nil, err
			}
		} else if mc.cfg != nil && mc.cfg.allowCleartextPasswords && err == ErrCleartextPassword {
			if err = mc.writeClearAuthPacket(); err != nil {
				mc.Close()
				return nil, err
			}
			if err = mc.readResultOK(); err != nil {
				mc.Close()
				return nil, err
			}
		} else {
			mc.Close()
			return nil, err
		}

	}

	// Get max allowed packet size
	maxap, err := mc.getSystemVar("max_allowed_packet")
	if err != nil {
		mc.Close()
		return nil, err
	}
	mc.maxPacketAllowed = stringToInt(maxap) - 1
	if mc.maxPacketAllowed < maxPacketSize {
		mc.maxWriteSize = mc.maxPacketAllowed
	}

	// Handle DSN Params
	err = mc.handleParams()
	if err != nil {
		mc.Close()
		return nil, err
	}

	return mc, nil
}

func init() {
	sql.Register("mysql", &MySQLDriver{})
}
<commit_msg>Add deadline for connection<commit_after>// Copyright 2012 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// The driver should be used via the database/sql package:
//
//  import "database/sql"
//  import _ "github.com/go-sql-driver/mysql"
//
//  db, err := sql.Open("mysql", "user:password@/dbname")
//
// See https://github.com/go-sql-driver/mysql#usage for details
package mysql

import (
	"database/sql"
	"database/sql/driver"
	"net"
	"time"
)

// This struct is exported to make the driver directly accessible.
// In general the driver is used via the database/sql package.
type MySQLDriver struct{}

// DialFunc is a function which can be used to establish the network connection.
// Custom dial functions must be registered with RegisterDial
type DialFunc func(addr string) (net.Conn, error)

var dials map[string]DialFunc

// RegisterDial registers a custom dial function. It can then be used by the
// network address mynet(addr), where mynet is the registered new network.
// addr is passed as a parameter to the dial function.
func RegisterDial(net string, dial DialFunc) {
	if dials == nil {
		dials = make(map[string]DialFunc)
	}
	dials[net] = dial
}

// Open new Connection.
// See https://github.com/go-sql-driver/mysql#dsn-data-source-name for how
// the DSN string is formated
func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
	var err error

	// New mysqlConn
	mc := &mysqlConn{
		maxPacketAllowed: maxPacketSize,
		maxWriteSize:     maxPacketSize - 1,
	}
	mc.cfg, err = parseDSN(dsn)
	if err != nil {
		return nil, err
	}

	// Connect to Server
	if dial, ok := dials[mc.cfg.net]; ok {
		mc.netConn, err = dial(mc.cfg.addr)
	} else {
		nd := net.Dialer{Timeout: mc.cfg.timeout}
		mc.netConn, err = nd.Dial(mc.cfg.net, mc.cfg.addr)
	}
	if err != nil {
		return nil, err
	}

	// Enable TCP Keepalives on TCP connections
	if tc, ok := mc.netConn.(*net.TCPConn); ok {
		timeout := time.Now().Add(mc.cfg.timeout)
		tc.SetReadDeadline(timeout)
		tc.SetDeadline(timeout)
		if err := tc.SetKeepAlive(true); err != nil {
			// Don't send COM_QUIT before handshake.
			mc.netConn.Close()
			mc.netConn = nil
			return nil, err
		}
	}

	mc.buf = newBuffer(mc.netConn)

	// Reading Handshake Initialization Packet
	cipher, err := mc.readInitPacket()
	if err != nil {
		mc.Close()
		return nil, err
	}

	// Send Client Authentication Packet
	if err = mc.writeAuthPacket(cipher); err != nil {
		mc.Close()
		return nil, err
	}

	// Read Result Packet
	err = mc.readResultOK()
	if err != nil {
		// Retry with old authentication method, if allowed
		if mc.cfg != nil && mc.cfg.allowOldPasswords && err == ErrOldPassword {
			if err = mc.writeOldAuthPacket(cipher); err != nil {
				mc.Close()
				return nil, err
			}
			if err = mc.readResultOK(); err != nil {
				mc.Close()
				return nil, err
			}
		} else if mc.cfg != nil && mc.cfg.allowCleartextPasswords && err == ErrCleartextPassword {
			if err = mc.writeClearAuthPacket(); err != nil {
				mc.Close()
				return nil, err
			}
			if err = mc.readResultOK(); err != nil {
				mc.Close()
				return nil, err
			}
		} else {
			mc.Close()
			return nil, err
		}

	}

	// Get max allowed packet size
	maxap, err := mc.getSystemVar("max_allowed_packet")
	if err != nil {
		mc.Close()
		return nil, err
	}
	mc.maxPacketAllowed = stringToInt(maxap) - 1
	if mc.maxPacketAllowed < maxPacketSize {
		mc.maxWriteSize = mc.maxPacketAllowed
	}

	// Handle DSN Params
	err = mc.handleParams()
	if err != nil {
		mc.Close()
		return nil, err
	}

	return mc, nil
}

func init() {
	sql.Register("mysql", &MySQLDriver{})
}
<|endoftext|>