<commit_before>#include<iostream>
#include<cstdio>
#include<sstream>
#include<fstream>
#include<iomanip>
#include<unistd.h>
#include<map>
#include<sys/types.h>
#include<sys/wait.h>
#include"utils.h"
#include"sandbox.h"
#include"config.h"
#include"testsuite.h"
#include"server_io.h"

using namespace std;

int compile(const submission& target, int boxid, int spBoxid);
void eval(submission &sub, int td, int boxid, int spBoxid);
void getExitStatus(submission &sub, int td);

int MAXPARNUM = 1;
int BOXOFFSET = 10;
bool AGGUPDATE = false;

int testsuite(submission &sub)
{
   system("rm -f ./testzone/*");
   const int testBoxid = BOXOFFSET + 0, spBoxid = BOXOFFSET + 1;
   sandboxInit(testBoxid);
   sandboxInit(spBoxid);
   int status = compile(sub, testBoxid, spBoxid);
   if(status != OK) return status;

   //anyway, only have batch judge right now
   map<pid_t, int> proc;
   int procnum = 0;
   int problem_id = sub.problem_id;
   int* time_limit = sub.time_limit;
   int* mem_limit = sub.mem_limit;

   for(int i = 0; i < sub.testdata_count; ++i){
      while(procnum >= MAXPARNUM){
         int status;
         pid_t cid = waitpid(-1, &status, 0);
         if(cid == -1){
            perror("[ERROR] in testsuite,  `waitpid()` failed :");
            return ER;
         }
         int td = proc[cid];
         eval(sub, td, BOXOFFSET + 10 + td, spBoxid);
         if(AGGUPDATE){
            sendResult(sub, OK, false);
         }
         //cerr << "td" << td << " : " << sub.verdict[td] << endl;
         sandboxDele(BOXOFFSET + 10 + td);
         --procnum;
      }

      if(procnum < MAXPARNUM){
         //batch judge
         ostringstream command;
         command << "batchjudge " << problem_id;
         command << " " << i;
         command << " " << BOXOFFSET + 10 + i;
         command << " " << time_limit[i];
         command << " " << mem_limit[i];
         command << " " << testBoxid;
         //
         pid_t pid = fork();
         if(pid == -1){
            perror("[ERROR] in testsuite, `fork()` failed :");
            return ER;
         }
         if(pid == 0){
            //child proc
            execl("/bin/sh", "sh", "-c", command.str().c_str(), NULL);
            perror("[ERROR] in testsuite, `execl()` failed :");
            exit(0);
         }
         proc[pid] = i;
         ++procnum;
      }
   }
   while(procnum >= 1){
      int status;
      pid_t cid = waitpid(-1, &status, 0);
      if(cid == -1){
         perror("[ERROR] in testsuite,  `waitpid()` failed :");
         return ER;
      }
      const int td = proc[cid];
      //sub.verdict[td] = eval(problem_id, td);
      eval(sub, td, BOXOFFSET + 10 + td, spBoxid);
      if(AGGUPDATE){
         sendResult(sub, OK, false);
      }
      //cerr << "td" << td << " : " << sub.verdict[td] << endl;
      sandboxDele(BOXOFFSET + 10 + td);
      --procnum;
   }
   //clear box-10
   sandboxDele(testBoxid);
   sandboxDele(spBoxid);
   return OK;
}

void setExitStatus(submission &sub, int td)
{
   ostringstream sout;
   sout << "./testzone/META" << td;
   ifstream fin(sout.str());
   string line;
   map<string,string> META;
   while(!fin.eof() && getline(fin, line)){
      for(int i = 0; i < line.size(); ++i)
         if(line[i] == ':')
            line[i] = ' ';
      istringstream in(line);
      string a, b;
      in >> a >> b;
      META[a] = b;
   }

   //mem_used
   sub.mem[td] = cast(META["max-rss"]).to<int>();
   //time_used
   sub.time[td] = 1000 * cast(META["time"]).to<double>();
   //verdict
   if(META["status"] == ""){
      sub.verdict[td] = OK;
   }else if(META["status"] == "TO"){
      sub.verdict[td] = TLE;
   }else if(META["status"] == "SG"){
      switch(cast(META["exitsig"]).to<int>()){
         case 11:
            sub.verdict[td] = MLE;
            break;
         default :
            sub.verdict[td] = RE;
      }
   }else if(META["status"] == "RE"){
      sub.verdict[td] = RE;
   }else{
      // "XX"
      sub.verdict[td] = ER;
   }
   return ;
}

void eval(submission &sub, int td, int boxid, int spBoxid)
{
   int problem_id = sub.problem_id;
   setExitStatus(sub, td);
   if(sub.verdict[td] != OK){
      return ;
   }

   if(sub.problem_type == 1){
      ostringstream sout;
      sout << "/tmp/box/" << spBoxid << "/box/sj.out ";
      string sjpath(sout.str());
      sout.str("");
      sout << "./testdata" << setfill('0') << setw(4) << problem_id << "/input" << setw(3) << td << ' ';
      string tdin(sout.str());
      sout.str("");
      sout << "./testdata" << setfill('0') << setw(4) << problem_id << "/output" << setw(3) << td << ' ';
      string tdout(sout.str());
      sout.str("");
      sout << "/tmp/box/" << boxid << "/box/output ";
      string ttout(sout.str());
      FILE* Pipe = popen((sjpath+ttout+tdin+tdout).c_str(), "r");
      int result = 1;
      fscanf(Pipe, "%d", &result);
      pclose(Pipe);
      if(result == 0)
         sub.verdict[td] = AC;
      else
         sub.verdict[td] = WA;
      return ;
   }

   int status = AC;
   string s, t;
   //solution output
   ostringstream sout;
   sout << "./testdata/" << setfill('0') << setw(4) << problem_id
      << "/output" << setw(3) << td;
   fstream tsol(sout.str());
   //user output
   sout.str("");
   sout << "/tmp/box/" << boxid << "/box/output";
   fstream mout(sout.str());
   while(1){
      s="";
      t="";
      getline(tsol,s);
      getline(mout,t);
      if(t==""&&s!=""){
         status = WA;
         break;
      }
      if(t!=""&&s==""){
         status = WA;
         break;
      }
      if(t==""&&s==""){
         break;
      }
      s.erase(s.find_last_not_of(" \n\r\t")+1);
      t.erase(t.find_last_not_of(" \n\r\t")+1);
      if(s!=t){
         status = WA;
         break;
      }
   }
   sub.verdict[td] = status;
   return ;
}

int compile(const submission& target, int boxid, int spBoxid)
{
   ostringstream sout;
   sout << "/tmp/box/" << boxid << "/box/";
   string boxdir(sout.str());

   ofstream fout;
   if(target.lang == "c++"){
      fout.open(boxdir + "main.cpp");
   }else{
      fout.open(boxdir + "main.c");
   }
   fout << target.code << flush;
   fout.close();

   if(target.problem_type == 2){
      sout.str("");
      sout << boxdir << "lib" << setw(4) << setfill('0') << target.problem_id << ".h";
      fout.open(sout.str());
      fout << target.interlib << flush;
      fout.close();
   }

   sout.str("");
   if(target.lang == "c++"){
      sout << "/usr/bin/g++ ./main.cpp -o ./main.out -O2 ";
   }else{
      sout << "/usr/bin/gcc ./main.c -o ./main.out -O2 -ansi ";
   }
   if(!target.std.empty()){
      sout << "-std=" << target.std << " ";
   }
   string comm(sout.str());
   sandboxOptions opt;
   opt.dirs.push_back("/etc/");
   opt.procs = 50;
   opt.preserve_env = true;
   opt.errout = "compile_error";
   opt.timeout = 60 * 1000;
   opt.meta = "./testzone/metacomp";

   sandboxExec(boxid, opt, comm);
   if(access((boxdir+"main.out").c_str(), F_OK) == -1){
      return CE;
   }

   if(target.problem_type == 1){
      sout.str("");
      sout << "/tmp/box/" << spBoxid << "/box/";
      string spBoxdir = sout.str();

      fout.open(spBoxdir + "sj.cpp");
      fout << target.sjcode << flush;
      fout.close();

      sout.str("");
      sout << "/usr/bin/g++ ./sj.cpp -o ./sj.out -O2 -std=c++11 ";

      opt.meta = "./testzone/metacompsj";

      sandboxExec(spBoxid, opt, sout.str());
      if(access((spBoxdir+"sj.out").c_str(), F_OK) == -1){
         return ER;
      }
   }

   return OK;
}
<commit_msg>jizzzzzzzzzzzzz<commit_after>#include<iostream>
#include<cstdio>
#include<sstream>
#include<fstream>
#include<iomanip>
#include<unistd.h>
#include<map>
#include<sys/types.h>
#include<sys/wait.h>
#include"utils.h"
#include"sandbox.h"
#include"config.h"
#include"testsuite.h"
#include"server_io.h"

using namespace std;

int compile(const submission& target, int boxid, int spBoxid);
void eval(submission &sub, int td, int boxid, int spBoxid);
void getExitStatus(submission &sub, int td);

int MAXPARNUM = 1;
int BOXOFFSET = 10;
bool AGGUPDATE = false;

int testsuite(submission &sub)
{
   system("rm -f ./testzone/*");
   const int testBoxid = BOXOFFSET + 0, spBoxid = BOXOFFSET + 1;
   sandboxInit(testBoxid);
   sandboxInit(spBoxid);
   int status = compile(sub, testBoxid, spBoxid);
   if(status != OK) return status;

   //anyway, only have batch judge right now
   map<pid_t, int> proc;
   int procnum = 0;
   int problem_id = sub.problem_id;
   int* time_limit = sub.time_limit;
   int* mem_limit = sub.mem_limit;

   for(int i = 0; i < sub.testdata_count; ++i){
      while(procnum >= MAXPARNUM){
         int status;
         pid_t cid = waitpid(-1, &status, 0);
         if(cid == -1){
            perror("[ERROR] in testsuite,  `waitpid()` failed :");
            return ER;
         }
         int td = proc[cid];
         eval(sub, td, BOXOFFSET + 10 + td, spBoxid);
         if(AGGUPDATE){
            sendResult(sub, OK, false);
         }
         //cerr << "td" << td << " : " << sub.verdict[td] << endl;
         sandboxDele(BOXOFFSET + 10 + td);
         --procnum;
      }

      if(procnum < MAXPARNUM){
         //batch judge
         ostringstream command;
         command << "batchjudge " << problem_id;
         command << " " << i;
         command << " " << BOXOFFSET + 10 + i;
         command << " " << time_limit[i];
         command << " " << mem_limit[i];
         command << " " << testBoxid;
         //
         pid_t pid = fork();
         if(pid == -1){
            perror("[ERROR] in testsuite, `fork()` failed :");
            return ER;
         }
         if(pid == 0){
            //child proc
            execl("/bin/sh", "sh", "-c", command.str().c_str(), NULL);
            perror("[ERROR] in testsuite, `execl()` failed :");
            exit(0);
         }
         proc[pid] = i;
         ++procnum;
      }
   }
   while(procnum >= 1){
      int status;
      pid_t cid = waitpid(-1, &status, 0);
      if(cid == -1){
         perror("[ERROR] in testsuite,  `waitpid()` failed :");
         return ER;
      }
      const int td = proc[cid];
      //sub.verdict[td] = eval(problem_id, td);
      eval(sub, td, BOXOFFSET + 10 + td, spBoxid);
      if(AGGUPDATE){
         sendResult(sub, OK, false);
      }
      //cerr << "td" << td << " : " << sub.verdict[td] << endl;
      sandboxDele(BOXOFFSET + 10 + td);
      --procnum;
   }
   //clear box-10
   sandboxDele(testBoxid);
   sandboxDele(spBoxid);
   return OK;
}

void setExitStatus(submission &sub, int td)
{
   ostringstream sout;
   sout << "./testzone/META" << td;
   ifstream fin(sout.str());
   string line;
   map<string,string> META;
   while(!fin.eof() && getline(fin, line)){
      for(int i = 0; i < line.size(); ++i)
         if(line[i] == ':')
            line[i] = ' ';
      istringstream in(line);
      string a, b;
      in >> a >> b;
      META[a] = b;
   }

   //mem_used
   sub.mem[td] = cast(META["max-rss"]).to<int>();
   //time_used
   sub.time[td] = 1000 * cast(META["time"]).to<double>();
   //verdict
   if(META["status"] == ""){
      sub.verdict[td] = OK;
   }else if(META["status"] == "TO"){
      sub.verdict[td] = TLE;
   }else if(META["status"] == "SG"){
      switch(cast(META["exitsig"]).to<int>()){
         case 11:
            sub.verdict[td] = MLE;
            break;
         default :
            sub.verdict[td] = RE;
      }
   }else if(META["status"] == "RE"){
      sub.verdict[td] = RE;
   }else{
      // "XX"
      sub.verdict[td] = ER;
   }
   return ;
}

void eval(submission &sub, int td, int boxid, int spBoxid)
{
   int problem_id = sub.problem_id;
   setExitStatus(sub, td);
   if(sub.verdict[td] != OK){
      return ;
   }

   if(sub.problem_type == 1){
      ostringstream sout;
      sout << "/tmp/box/" << spBoxid << "/box/sj.out ";
      string sjpath(sout.str());
      sout.str("");
      sout << "./testdata" << setfill('0') << setw(4) << problem_id << "/input" << setw(3) << td << ' ';
      string tdin(sout.str());
      sout.str("");
      sout << "./testdata" << setfill('0') << setw(4) << problem_id << "/output" << setw(3) << td << ' ';
      string tdout(sout.str());
      sout.str("");
      sout << "/tmp/box/" << boxid << "/box/output ";
      string ttout(sout.str());
      FILE* Pipe = popen((sjpath+ttout+tdin+tdout).c_str(), "r");
      int result = 1;
      fscanf(Pipe, "%d", &result);
      pclose(Pipe);
      if(result == 0)
         sub.verdict[td] = AC;
      else
         sub.verdict[td] = WA;
      return ;
   }

   int status = AC;
   //solution output
   ostringstream sout;
   sout << "./testdata/" << setfill('0') << setw(4) << problem_id
      << "/output" << setw(3) << td;
   fstream tsol(sout.str());
   //user output
   sout.str("");
   sout << "/tmp/box/" << boxid << "/box/output";
   fstream mout(sout.str());
   while(true){
      if(tsol.eof() != mout.eof()){
         status = WA;
         break;
      }
      if(tsol.eof() && mout.eof()){
         break;
      }
      string s, t;
      getline(tsol,s);
      getline(mout,t);
      s.erase(s.find_last_not_of(" \n\r\t")+1);
      t.erase(t.find_last_not_of(" \n\r\t")+1);
      if(s != t){
         status = WA;
         break;
      }
   }
   sub.verdict[td] = status;
   return ;
}

int compile(const submission& target, int boxid, int spBoxid)
{
   ostringstream sout;
   sout << "/tmp/box/" << boxid << "/box/";
   string boxdir(sout.str());

   ofstream fout;
   if(target.lang == "c++"){
      fout.open(boxdir + "main.cpp");
   }else{
      fout.open(boxdir + "main.c");
   }
   fout << target.code << flush;
   fout.close();

   if(target.problem_type == 2){
      sout.str("");
      sout << boxdir << "lib" << setw(4) << setfill('0') << target.problem_id << ".h";
      fout.open(sout.str());
      fout << target.interlib << flush;
      fout.close();
   }

   sout.str("");
   if(target.lang == "c++"){
      sout << "/usr/bin/g++ ./main.cpp -o ./main.out -O2 ";
   }else{
      sout << "/usr/bin/gcc ./main.c -o ./main.out -O2 -ansi ";
   }
   if(!target.std.empty()){
      sout << "-std=" << target.std << " ";
   }
   string comm(sout.str());
   sandboxOptions opt;
   opt.dirs.push_back("/etc/");
   opt.procs = 50;
   opt.preserve_env = true;
   opt.errout = "compile_error";
   opt.timeout = 60 * 1000;
   opt.meta = "./testzone/metacomp";

   sandboxExec(boxid, opt, comm);
   if(access((boxdir+"main.out").c_str(), F_OK) == -1){
      return CE;
   }

   if(target.problem_type == 1){
      sout.str("");
      sout << "/tmp/box/" << spBoxid << "/box/";
      string spBoxdir = sout.str();

      fout.open(spBoxdir + "sj.cpp");
      fout << target.sjcode << flush;
      fout.close();

      sout.str("");
      sout << "/usr/bin/g++ ./sj.cpp -o ./sj.out -O2 -std=c++11 ";

      opt.meta = "./testzone/metacompsj";

      sandboxExec(spBoxid, opt, sout.str());
      if(access((spBoxdir+"sj.out").c_str(), F_OK) == -1){
         return ER;
      }
   }

   return OK;
}
<|endoftext|>


<commit_before>/*
* Client Key Exchange Message
* (C) 2004-2010 Jack Lloyd
*
* Released under the terms of the Botan license
*/

#include <botan/internal/tls_messages.h>
#include <botan/internal/tls_reader.h>
#include <botan/internal/tls_extensions.h>
#include <botan/pubkey.h>
#include <botan/dh.h>
#include <botan/ecdh.h>
#include <botan/rsa.h>
#include <botan/rng.h>
#include <botan/loadstor.h>
#include <memory>

namespace Botan {

namespace TLS {

namespace {

SecureVector<byte> strip_leading_zeros(const MemoryRegion<byte>& input)
   {
   size_t leading_zeros = 0;

   for(size_t i = 0; i != input.size(); ++i)
      {
      if(input[i] != 0)
         break;
      ++leading_zeros;
      }

   SecureVector<byte> output(&input[leading_zeros],
                             input.size() - leading_zeros);
   return output;
   }

}

/*
* Create a new Client Key Exchange message
*/
Client_Key_Exchange::Client_Key_Exchange(Record_Writer& writer,
                                         Handshake_State* state,
                                         const std::vector<X509_Certificate>& peer_certs,
                                         RandomNumberGenerator& rng)
   {
   if(state->server_kex)
      {
      TLS_Data_Reader reader(state->server_kex->params());

      if(state->suite.kex_algo() == "DH")
         {
         BigInt p = BigInt::decode(reader.get_range<byte>(2, 1, 65535));
         BigInt g = BigInt::decode(reader.get_range<byte>(2, 1, 65535));
         BigInt Y = BigInt::decode(reader.get_range<byte>(2, 1, 65535));

         if(reader.remaining_bytes())
            throw Decoding_Error("Bad params size for DH key exchange");

         DL_Group group(p, g);

         if(!group.verify_group(rng, true))
            throw Internal_Error("DH group failed validation, possible attack");

         DH_PublicKey counterparty_key(group, Y);

         // FIXME Check that public key is residue?

         DH_PrivateKey priv_key(rng, group);

         PK_Key_Agreement ka(priv_key, "Raw");

         pre_master = strip_leading_zeros(
            ka.derive_key(0, counterparty_key.public_value()).bits_of());

         append_tls_length_value(key_material, priv_key.public_value(), 2);
         }
      else if(state->suite.kex_algo() == "ECDH")
         {
         const byte curve_type = reader.get_byte();

         if(curve_type != 3)
            throw Decoding_Error("Server sent non-named ECC curve");

         const u16bit curve_id = reader.get_u16bit();

         const std::string name = Supported_Elliptic_Curves::curve_id_to_name(curve_id);

         if(name == "")
            throw Decoding_Error("Server sent unknown named curve " + to_string(curve_id));

         EC_Group group(name);

         MemoryVector<byte> ecdh_key = reader.get_range<byte>(1, 1, 255);

         ECDH_PublicKey counterparty_key(group, OS2ECP(ecdh_key, group.get_curve()));

         ECDH_PrivateKey priv_key(rng, group);

         PK_Key_Agreement ka(priv_key, "Raw");

         pre_master = strip_leading_zeros(
            ka.derive_key(0, counterparty_key.public_value()).bits_of());

         append_tls_length_value(key_material, priv_key.public_value(), 1);
         }
      else
         throw Internal_Error("Server key exchange type " + state->suite.kex_algo() +
                              " not known");
      }
   else
      {
      // No server key exchange msg better mean a RSA key in the cert

      std::auto_ptr<Public_Key> pub_key(peer_certs[0].subject_public_key());

      if(peer_certs.empty())
         throw Internal_Error("No certificate and no server key exchange");

      if(const RSA_PublicKey* rsa_pub = dynamic_cast<const RSA_PublicKey*>(pub_key.get()))
         {
         const Protocol_Version pref_version = state->client_hello->version();

         pre_master = rng.random_vec(48);
         pre_master[0] = pref_version.major_version();
         pre_master[1] = pref_version.minor_version();

         PK_Encryptor_EME encryptor(*rsa_pub, "PKCS1v15");

         MemoryVector<byte> encrypted_key = encryptor.encrypt(pre_master, rng);

         if(state->version == Protocol_Version::SSL_V3)
            key_material = encrypted_key; // no length field
         else
            append_tls_length_value(key_material, encrypted_key, 2);
         }
      else
         throw TLS_Exception(HANDSHAKE_FAILURE,
                             "Expected a RSA key in server cert but got " +
                             pub_key->algo_name());
      }

   send(writer, state->hash);
   }

/*
* Read a Client Key Exchange message
*/
Client_Key_Exchange::Client_Key_Exchange(const MemoryRegion<byte>& contents,
                                         const Ciphersuite& suite,
                                         Protocol_Version using_version)
   {
   if(suite.kex_algo() == "" && using_version == Protocol_Version::SSL_V3)
      key_material = contents;
   else
      {
      TLS_Data_Reader reader(contents);

      if(suite.kex_algo() == "" || suite.kex_algo() == "DH")
         key_material = reader.get_range<byte>(2, 0, 65535);
      else if(suite.kex_algo() == "ECDH")
         key_material = reader.get_range<byte>(1, 1, 255);
      else
         throw Internal_Error("Unknown client key exch type " + suite.kex_algo());
      }
   }

/*
* Return the pre_master_secret
*/
SecureVector<byte>
Client_Key_Exchange::pre_master_secret(RandomNumberGenerator& rng,
                                       const Private_Key* priv_key,
                                       Protocol_Version client_version)
   {

   if(const DH_PrivateKey* dh_priv = dynamic_cast<const DH_PrivateKey*>(priv_key))
      {
      try {
         PK_Key_Agreement ka(*dh_priv, "Raw");

         pre_master = strip_leading_zeros(ka.derive_key(0, key_material).bits_of());
      }
      catch(...)
         {
         /*
         * Something failed in the DH computation. To avoid possible
         * timing attacks, randomize the pre-master output and carry
         * on, allowing the protocol to fail later in the finished
         * checks.
         */
         pre_master = rng.random_vec(dh_priv->public_value().size());
         }

      return pre_master;
      }
   else if(const RSA_PrivateKey* rsa_priv = dynamic_cast<const RSA_PrivateKey*>(priv_key))
      {
      PK_Decryptor_EME decryptor(*rsa_priv, "PKCS1v15");

      try {
         pre_master = decryptor.decrypt(key_material);

         if(pre_master.size() != 48 ||
            client_version.major_version() != pre_master[0] ||
            client_version.minor_version() != pre_master[1])
            {
            throw Decoding_Error("Client_Key_Exchange: Secret corrupted");
            }
      }
      catch(...)
         {
         pre_master = rng.random_vec(48);
         pre_master[0] = client_version.major_version();
         pre_master[1] = client_version.minor_version();
         }

      return pre_master;
      }
   else
      throw Invalid_Argument("Client_Key_Exchange: Bad key for decrypt");
   }

}

}
<commit_msg>Read ECDH client key exchange messages<commit_after>/*
* Client Key Exchange Message
* (C) 2004-2010 Jack Lloyd
*
* Released under the terms of the Botan license
*/

#include <botan/internal/tls_messages.h>
#include <botan/internal/tls_reader.h>
#include <botan/internal/tls_extensions.h>
#include <botan/pubkey.h>
#include <botan/dh.h>
#include <botan/ecdh.h>
#include <botan/rsa.h>
#include <botan/rng.h>
#include <botan/loadstor.h>
#include <memory>

namespace Botan {

namespace TLS {

namespace {

SecureVector<byte> strip_leading_zeros(const MemoryRegion<byte>& input)
   {
   size_t leading_zeros = 0;

   for(size_t i = 0; i != input.size(); ++i)
      {
      if(input[i] != 0)
         break;
      ++leading_zeros;
      }

   SecureVector<byte> output(&input[leading_zeros],
                             input.size() - leading_zeros);
   return output;
   }

}

/*
* Create a new Client Key Exchange message
*/
Client_Key_Exchange::Client_Key_Exchange(Record_Writer& writer,
                                         Handshake_State* state,
                                         const std::vector<X509_Certificate>& peer_certs,
                                         RandomNumberGenerator& rng)
   {
   if(state->server_kex)
      {
      TLS_Data_Reader reader(state->server_kex->params());

      if(state->suite.kex_algo() == "DH")
         {
         BigInt p = BigInt::decode(reader.get_range<byte>(2, 1, 65535));
         BigInt g = BigInt::decode(reader.get_range<byte>(2, 1, 65535));
         BigInt Y = BigInt::decode(reader.get_range<byte>(2, 1, 65535));

         if(reader.remaining_bytes())
            throw Decoding_Error("Bad params size for DH key exchange");

         DL_Group group(p, g);

         if(!group.verify_group(rng, true))
            throw Internal_Error("DH group failed validation, possible attack");

         DH_PublicKey counterparty_key(group, Y);

         // FIXME Check that public key is residue?

         DH_PrivateKey priv_key(rng, group);

         PK_Key_Agreement ka(priv_key, "Raw");

         pre_master = strip_leading_zeros(
            ka.derive_key(0, counterparty_key.public_value()).bits_of());

         append_tls_length_value(key_material, priv_key.public_value(), 2);
         }
      else if(state->suite.kex_algo() == "ECDH")
         {
         const byte curve_type = reader.get_byte();

         if(curve_type != 3)
            throw Decoding_Error("Server sent non-named ECC curve");

         const u16bit curve_id = reader.get_u16bit();

         const std::string name = Supported_Elliptic_Curves::curve_id_to_name(curve_id);

         if(name == "")
            throw Decoding_Error("Server sent unknown named curve " + to_string(curve_id));

         EC_Group group(name);

         MemoryVector<byte> ecdh_key = reader.get_range<byte>(1, 1, 255);

         ECDH_PublicKey counterparty_key(group, OS2ECP(ecdh_key, group.get_curve()));

         ECDH_PrivateKey priv_key(rng, group);

         PK_Key_Agreement ka(priv_key, "Raw");

         pre_master = strip_leading_zeros(
            ka.derive_key(0, counterparty_key.public_value()).bits_of());

         append_tls_length_value(key_material, priv_key.public_value(), 1);
         }
      else
         throw Internal_Error("Server key exchange type " + state->suite.kex_algo() +
                              " not known");
      }
   else
      {
      // No server key exchange msg better mean a RSA key in the cert

      std::auto_ptr<Public_Key> pub_key(peer_certs[0].subject_public_key());

      if(peer_certs.empty())
         throw Internal_Error("No certificate and no server key exchange");

      if(const RSA_PublicKey* rsa_pub = dynamic_cast<const RSA_PublicKey*>(pub_key.get()))
         {
         const Protocol_Version pref_version = state->client_hello->version();

         pre_master = rng.random_vec(48);
         pre_master[0] = pref_version.major_version();
         pre_master[1] = pref_version.minor_version();

         PK_Encryptor_EME encryptor(*rsa_pub, "PKCS1v15");

         MemoryVector<byte> encrypted_key = encryptor.encrypt(pre_master, rng);

         if(state->version == Protocol_Version::SSL_V3)
            key_material = encrypted_key; // no length field
         else
            append_tls_length_value(key_material, encrypted_key, 2);
         }
      else
         throw TLS_Exception(HANDSHAKE_FAILURE,
                             "Expected a RSA key in server cert but got " +
                             pub_key->algo_name());
      }

   send(writer, state->hash);
   }

/*
* Read a Client Key Exchange message
*/
Client_Key_Exchange::Client_Key_Exchange(const MemoryRegion<byte>& contents,
                                         const Ciphersuite& suite,
                                         Protocol_Version using_version)
   {
   if(suite.kex_algo() == "" && using_version == Protocol_Version::SSL_V3)
      key_material = contents;
   else
      {
      TLS_Data_Reader reader(contents);

      if(suite.kex_algo() == "" || suite.kex_algo() == "DH")
         key_material = reader.get_range<byte>(2, 0, 65535);
      else if(suite.kex_algo() == "ECDH")
         key_material = reader.get_range<byte>(1, 1, 255);
      else
         throw Internal_Error("Unknown client key exch type " + suite.kex_algo());
      }
   }

/*
* Return the pre_master_secret
*/
SecureVector<byte>
Client_Key_Exchange::pre_master_secret(RandomNumberGenerator& rng,
                                       const Private_Key* priv_key,
                                       Protocol_Version client_version)
   {
   if(const RSA_PrivateKey* rsa = dynamic_cast<const RSA_PrivateKey*>(priv_key))
      {
      PK_Decryptor_EME decryptor(*rsa, "PKCS1v15");

      try {
         pre_master = decryptor.decrypt(key_material);

         if(pre_master.size() != 48 ||
            client_version.major_version() != pre_master[0] ||
            client_version.minor_version() != pre_master[1])
            {
            throw Decoding_Error("Client_Key_Exchange: Secret corrupted");
            }
      }
      catch(...)
         {
         pre_master = rng.random_vec(48);
         pre_master[0] = client_version.major_version();
         pre_master[1] = client_version.minor_version();
         }

      return pre_master;
      }

   // DH or ECDH
   if(const PK_Key_Agreement_Key* dh = dynamic_cast<const PK_Key_Agreement_Key*>(priv_key))
      {
      try {
         PK_Key_Agreement ka(*dh, "Raw");

         pre_master = strip_leading_zeros(ka.derive_key(0, key_material).bits_of());
      }
      catch(...)
         {
         /*
         * Something failed in the DH computation. To avoid possible
         * timing attacks, randomize the pre-master output and carry
         * on, allowing the protocol to fail later in the finished
         * checks.
         */
         pre_master = rng.random_vec(dh->public_value().size());
         }

      return pre_master;
      }

   throw Invalid_Argument("Client_Key_Exchange: Unknown key type " + priv_key->algo_name());
   }

}

}
<|endoftext|>


<commit_before>#pragma once
#include "../common.hpp"
#include <Prelude/HashTable.hpp>
#include <Prelude/Vector.hpp>
#include <Prelude/FastList.hpp>
#include "../generic/simple-set.hpp"
#include "../collector.hpp"
#include "nodes.hpp"

namespace Mirb
{
	class Block;

	namespace CodeGen
	{
		class Block;
	};

	namespace Tree
	{
		class Chunk
		{
			public:
				static const size_t main_size = 512;
				static const size_t block_size = 256;
				static const size_t allocation_limit = 128;

				static Chunk *create(size_t bytes);

				uint8_t *current;
				uint8_t *end;

				ListEntry<Chunk> entry;

				void *allocate(size_t bytes);
		};

		class Fragment:
			public WithReferenceProvider<Fragment>
		{
			public:
				static const bool can_free = false;

				Fragment(Fragment *parent, size_t chunk_size);

				Vector<Fragment *, Fragment> fragments;

				void *allocate(size_t bytes);
				void *reallocate(void *memory, size_t old_size, size_t new_size);
				void free(void *memory)
				{
				}

			private:
				size_t chunk_size;

				Chunk *current;

				FastList<Chunk> chunks;
		};

		struct Node;

		struct SuperNode;

		class Variable
		{
			public:
				enum Type
				{
					Temporary,
					Local,
					Heap,
					Types
				};

				Type type;
				Scope *owner; // Only valid for heap variables

				Variable(Type type)
					: type(type)
					#ifdef DEBUG
						, group(0)
					#endif
				{
				}

				size_t loc;

				#ifdef DEBUG
					Variable *group;
				#endif
		};

		class NamedVariable:
			public Variable
		{
			public:
				NamedVariable(Type type) : Variable(type) {}

				Symbol *name;
				NamedVariable *next;
		};

		class Parameter:
			public NamedVariable
		{
			public:
				Parameter(Type type) : NamedVariable(type) {}

				ListEntry<Parameter> parameter_entry;
		};

		class VariableMapFunctions:
			public HashTableFunctions<Symbol *, NamedVariable *, Fragment>
		{
			public:
				static bool compare_key_value(Symbol *key, NamedVariable *value)
				{
					return key == value->name;
				}

				static Symbol *get_key(NamedVariable *value)
				{
					return value->name;
				}

				static NamedVariable *get_value_next(NamedVariable *value)
				{
					return value->next;
				}

				static void set_value_next(NamedVariable *value, NamedVariable *next)
				{
					value->next = next;
				}

				static size_t hash_key(Symbol *key)
				{
					return (size_t)key;
				}

				static bool valid_key(Symbol *key)
				{
					return true;
				}
		};

		typedef HashTable<Symbol *, NamedVariable *, VariableMapFunctions, Fragment> VariableMap;

		class Scope:
			public PinnedHeader
		{
			public:
				enum Type
				{
					Top,
					Method,
					Class,
					Module,
					Closure
				};

				Scope(Document *document, Fragment *fragment, Scope *parent, Type type);

				Document *document;
				Block *final;
				Fragment *fragment;
				Type type;
				Scope *parent; // The scope enclosing this one
				Scope *owner; // The first parent that isn't a closure. This field can point to itself.
				Node *group;

				/*
				 * Break related fields
				 */
				size_t break_id; // Which of the parent's break targets this block belongs to.
				size_t break_targets; // Number of child blocks that can raise a break exception.
				var_t break_dst; // The variable in the parent that the break will override.

				static const size_t no_break_id = (size_t)-1;

				/*
				 * Exception related fields
				 */
				bool require_exceptions;

				/*
				 * Variable related fields
				 */
				VariableMap variables; // A hash of the variables in this scope.
				size_t heap_vars; // The number of the variables that must be stored on a heap scope.
				Parameter *block_parameter; // Pointer to a named or unnamed block variable.
				Vector<Scope *, Fragment> referenced_scopes; // A list of all the scopes this scope requires.

				Vector<Scope *, Fragment> children; // A list of all the immediate children. To keep them alive...

				Vector<Variable *, Fragment> variable_list; // A list of all variables in this scope.

				CountedList<Parameter, Parameter, &Parameter::parameter_entry> parameters;

				Vector<Scope *, Fragment> zsupers;

				template<typename F> void mark(F mark)
				{
					mark(final);
					mark(document);

					referenced_scopes.mark_content(mark);
					children.mark_content(mark);
					zsupers.mark_content(mark);

					if(parent)
						mark(parent);

					mark(owner);
				}

				template<class T> T *alloc_var(Variable::Type type = Variable::Local)
				{
					T *result = new (fragment) T(type);

					result->loc = variable_list.size();

					result->type = type;

					variable_list.push(result);

					return result;
				}

				template<class T> T *define(Symbol *name)
				{
					T *result = (T *)variables.get(name);

					if(result)
						return result;

					result = alloc_var<T>();

					result->name = name;

					variables.set(name, result);

					return result;
				}

				template<class T> T *get_var(Symbol *name)
				{
					Scope *defined_scope = defined(name, true);

					if(defined_scope == this)
					{
						return variables.get(name);
					}
					else if(defined_scope)
					{
						auto result = defined_scope->variables.get(name);

						require_var(defined_scope, result);

						return result;
					}
					else
						return define<T>(name);
				}

				void require_args(Scope *owner);
				void require_scope(Scope *scope);
				void require_var(Scope *owner, Variable *var);

				Scope *defined(Symbol *name, bool recursive);
		};
	};
};

void *operator new(size_t bytes, Mirb::Tree::Fragment *fragment) throw();
void operator delete(void *, Mirb::Tree::Fragment *fragment) throw();
void *operator new[](size_t bytes, Mirb::Tree::Fragment *fragment) throw();
void operator delete[](void *, Mirb::Tree::Fragment *fragment) throw();
<commit_msg>Test suite passes!<commit_after>#pragma once
#include "../common.hpp"
#include <Prelude/HashTable.hpp>
#include <Prelude/Vector.hpp>
#include <Prelude/FastList.hpp>
#include "../generic/simple-set.hpp"
#include "../collector.hpp"
#include "nodes.hpp"

namespace Mirb
{
	class Block;

	namespace CodeGen
	{
		class Block;
	};

	namespace Tree
	{
		class Chunk
		{
			public:
				static const size_t main_size = 512;
				static const size_t block_size = 256;
				static const size_t allocation_limit = 128;

				static Chunk *create(size_t bytes);

				uint8_t *current;
				uint8_t *end;

				ListEntry<Chunk> entry;

				void *allocate(size_t bytes);
		};

		class Fragment:
			public WithReferenceProvider<Fragment>
		{
			public:
				static const bool can_free = false;

				Fragment(Fragment *parent, size_t chunk_size);

				Vector<Fragment *, Fragment> fragments;

				void *allocate(size_t bytes);
				void *reallocate(void *memory, size_t old_size, size_t new_size);
				void free(void *memory)
				{
				}

			private:
				size_t chunk_size;

				Chunk *current;

				FastList<Chunk> chunks;
		};

		struct Node;

		struct SuperNode;

		class Variable
		{
			public:
				enum Type
				{
					Temporary,
					Local,
					Heap,
					Types
				};

				Type type;
				Scope *owner; // Only valid for heap variables

				Variable(Type type)
					: type(type)
					#ifdef DEBUG
						, group(0)
					#endif
				{
				}

				size_t loc;

				#ifdef DEBUG
					Variable *group;
				#endif
		};

		class NamedVariable:
			public Variable
		{
			public:
				NamedVariable(Type type) : Variable(type) {}

				Symbol *name;
				NamedVariable *next;
		};

		class Parameter:
			public NamedVariable
		{
			public:
				Parameter(Type type) : NamedVariable(type) {}

				ListEntry<Parameter> parameter_entry;
		};

		class VariableMapFunctions:
			public HashTableFunctions<Symbol *, NamedVariable *, Fragment>
		{
			public:
				static bool compare_key_value(Symbol *key, NamedVariable *value)
				{
					return key == value->name;
				}

				static Symbol *get_key(NamedVariable *value)
				{
					return value->name;
				}

				static NamedVariable *get_value_next(NamedVariable *value)
				{
					return value->next;
				}

				static void set_value_next(NamedVariable *value, NamedVariable *next)
				{
					value->next = next;
				}

				static size_t hash_key(Symbol *key)
				{
					return (size_t)key;
				}

				static bool valid_key(Symbol *key)
				{
					return true;
				}
		};

		typedef HashTable<Symbol *, NamedVariable *, VariableMapFunctions, Fragment> VariableMap;

		class Scope:
			public PinnedHeader
		{
			public:
				enum Type
				{
					Top,
					Method,
					Class,
					Module,
					Closure
				};

				Scope(Document *document, Fragment *fragment, Scope *parent, Type type);

				Document *document;
				Block *final;
				Fragment *fragment;
				Type type;
				Scope *parent; // The scope enclosing this one
				Scope *owner; // The first parent that isn't a closure. This field can point to itself.
				Node *group;

				/*
				 * Break related fields
				 */
				size_t break_id; // Which of the parent's break targets this block belongs to.
				size_t break_targets; // Number of child blocks that can raise a break exception.
				var_t break_dst; // The variable in the parent that the break will override.

				static const size_t no_break_id = (size_t)-1;

				/*
				 * Exception related fields
				 */
				bool require_exceptions;

				/*
				 * Variable related fields
				 */
				VariableMap variables; // A hash of the variables in this scope.
				size_t heap_vars; // The number of the variables that must be stored on a heap scope.
				Parameter *block_parameter; // Pointer to a named or unnamed block variable.
				Vector<Scope *, Fragment> referenced_scopes; // A list of all the scopes this scope requires.

				Vector<Scope *, Fragment> children; // A list of all the immediate children. To keep them alive...

				Vector<Variable *, Fragment> variable_list; // A list of all variables in this scope.

				CountedList<Parameter, Parameter, &Parameter::parameter_entry> parameters;

				Vector<Scope *, Fragment> zsupers;

				template<typename F> void mark(F mark)
				{
					if(final)
						mark(final);

					mark(document);

					referenced_scopes.mark_content(mark);
					children.mark_content(mark);
					zsupers.mark_content(mark);

					if(parent)
						mark(parent);

					mark(owner);
				}

				template<class T> T *alloc_var(Variable::Type type = Variable::Local)
				{
					T *result = new (fragment) T(type);

					result->loc = variable_list.size();

					result->type = type;

					variable_list.push(result);

					return result;
				}

				template<class T> T *define(Symbol *name)
				{
					T *result = (T *)variables.get(name);

					if(result)
						return result;

					result = alloc_var<T>();

					result->name = name;

					variables.set(name, result);

					return result;
				}

				template<class T> T *get_var(Symbol *name)
				{
					Scope *defined_scope = defined(name, true);

					if(defined_scope == this)
					{
						return variables.get(name);
					}
					else if(defined_scope)
					{
						auto result = defined_scope->variables.get(name);

						require_var(defined_scope, result);

						return result;
					}
					else
						return define<T>(name);
				}

				void require_args(Scope *owner);
				void require_scope(Scope *scope);
				void require_var(Scope *owner, Variable *var);

				Scope *defined(Symbol *name, bool recursive);
		};
	};
};

void *operator new(size_t bytes, Mirb::Tree::Fragment *fragment) throw();
void operator delete(void *, Mirb::Tree::Fragment *fragment) throw();
void *operator new[](size_t bytes, Mirb::Tree::Fragment *fragment) throw();
void operator delete[](void *, Mirb::Tree::Fragment *fragment) throw();
<|endoftext|>


<commit_before>/*** tws_strat.cpp -- TWS strategy module
 *
 * Copyright (C) 2012 Ruediger Meier
 *
 * Author:  Ruediger Meier <sweet_f_a@gmx.de>
 *
 * This file is part of twstools.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the author nor the names of any contributors
 *    may be used to endorse or promote products derived from this
 *    software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR "AS IS" AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ***/

#include "tws_strat.h"
#include "tws_query.h"
#include "tws_meta.h"
#include "twsdo.h"
#include "debug.h"
#include <assert.h>

buy_sell_oid::buy_sell_oid() :
	sell_oid(0),
	buy_oid(0)
{
}

Strat::Strat( TwsDL& _twsdo) :
	twsdo(_twsdo)
{
}

Strat::~Strat()
{
}

void Strat::adjustOrders()
{
	static int fuck = -1;
	fuck++;
	if( fuck <= 30 || twsdo.p_orders.size() > 0 ) {
		return;
	}

	DEBUG_PRINTF( "Adjust orders." );
	PlaceOrder pO;
	int i;

	const MktDataTodo &mtodo = twsdo.workTodo->getMktDataTodo();
	for( int i=0; i < mtodo.mktDataRequests.size(); i++ ) {
		pO.contract = mtodo.mktDataRequests[i].ibContract;
		pO.order.orderType = "LMT";
		pO.order.action = "BUY";
		pO.order.lmtPrice = twsdo.quotes->at(i).val[IB::BID] - 0.1;
		pO.order.totalQuantity = pO.contract.secType == "CASH" ? 25000 : 1;
		twsdo.workTodo->placeOrderTodo()->add(pO);
	}
	DEBUG_PRINTF( "Adjust orders. %d",
		twsdo.workTodo->placeOrderTodo()->countLeft());
}<commit_msg>Strat, prepare order id handling<commit_after>/*** tws_strat.cpp -- TWS strategy module
 *
 * Copyright (C) 2012 Ruediger Meier
 *
 * Author:  Ruediger Meier <sweet_f_a@gmx.de>
 *
 * This file is part of twstools.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the author nor the names of any contributors
 *    may be used to endorse or promote products derived from this
 *    software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR "AS IS" AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ***/

#include "tws_strat.h"
#include "tws_query.h"
#include "tws_meta.h"
#include "twsdo.h"
#include "debug.h"
#include <assert.h>

buy_sell_oid::buy_sell_oid() :
	sell_oid(0),
	buy_oid(0)
{
}

Strat::Strat( TwsDL& _twsdo) :
	twsdo(_twsdo)
{
}

Strat::~Strat()
{
}

void Strat::adjustOrders()
{
	static int fuck = -1;
	fuck++;
	if( fuck <= 30 || twsdo.p_orders.size() > 0 ) {
		return;
	}

	DEBUG_PRINTF( "strat, adjust orders" );

	const MktDataTodo &mtodo = twsdo.workTodo->getMktDataTodo();
	for( int i=0; i < mtodo.mktDataRequests.size(); i++ ) {
		PlaceOrder pO;
		pO.contract = mtodo.mktDataRequests[i].ibContract;
		pO.order.orderType = "LMT";
		pO.order.totalQuantity = pO.contract.secType == "CASH" ? 25000 : 1;
		const char *symbol = pO.contract.symbol.c_str();

		/* here we also initialize zero order ids */
		buy_sell_oid &oids = map_data_order[i];
		if( twsdo.p_orders.find(oids.buy_oid) == twsdo.p_orders.end() ) {
			/* new buy order */
			DEBUG_PRINTF( "strat, new buy order %s", symbol );
		} else {
			/* modify buy order */
			DEBUG_PRINTF( "strat, modify buy order %s", symbol );
		}
		if( twsdo.p_orders.find(oids.sell_oid) == twsdo.p_orders.end() ) {
			/* new sell order */
			DEBUG_PRINTF( "strat, new sell order %s", symbol );
		} else {
			/* modify sell order */
			DEBUG_PRINTF( "strat, modify sell order %s", symbol );
		}

		pO.order.action = "BUY";
		pO.order.lmtPrice = twsdo.quotes->at(i).val[IB::BID] - 0.1;
		twsdo.workTodo->placeOrderTodo()->add(pO);
	}
	DEBUG_PRINTF( "strat, place/modify %d orders",
		twsdo.workTodo->placeOrderTodo()->countLeft());
	assert( mtodo.mktDataRequests.size() == map_data_order.size() );
}<|endoftext|>


<commit_before>// SuperCard Pro flux file format:
//  http://www.cbmstuff.com/downloads/scp/scp_image_specs.txt

#include "SAMdisk.h"
#include "DemandDisk.h"

enum
{
    FLAG_INDEX = 1 << 0,    // set for index-synchronised, clear if not
    FLAG_TPI = 1 << 1,  // set for 96tpi, clear for 48tpi
    FLAG_RPM = 1 << 2,  // set for 360rpm, clear for 300rpm
    FLAG_TYPE = 1 << 3, // set for normalised flux, clear for preservation quality
    FLAG_MODE = 1 << 4, // set for read/write image, clear for read-only
    FLAG_FOOTER = 1 << 5    // set if footer block present, clear if absent
};

struct SCP_FILE_HEADER
{
    char signature[3];      // SCP
    uint8_t revision;       // (version << 4) | revision, 0x39 = v3.9
    uint8_t disk_type;
    uint8_t revolutions;    // number of stored revolutions
    uint8_t start_track;
    uint8_t end_track;
    uint8_t flags;
    uint8_t bitcell_width;  // 0 = 16 bits, non-zero = number of bits
    uint8_t heads;          // 0 = both heads, 1 = head 0 only, 2 = head 1 only
    uint8_t reserved;       // should be zero?
    uint32_t checksum;      // 32-bit checksum from after header to EOF (unless FLAG_MODE is set)
};

struct SCP_FILE_FOOTER
{
    uint32_t manufacturer_offset;
    uint32_t model_offset;
    uint32_t serial_offset;
    uint32_t creator_offset;
    uint32_t application_offset;
    uint32_t comments_offset;
    uint64_t creation_time;
    uint64_t modification_time;
    uint8_t application_version;
    uint8_t hardware_version;
    uint8_t firmware_version;
    uint8_t format_revision;
    char sig[4];
};

struct TRACK_DATA_HEADER
{
    char signature[3];      // TRK
    uint8_t tracknr;
};

std::string FooterString(MemFile& file, uint32_t offset)
{
    uint16_t len = 0;
    if (!offset || !file.seek(static_cast<int>(offset)) || !file.read(&len, sizeof(len)))
        return "";

    len = util::letoh(len);
    if (static_cast<int>(offset + len) >= file.size())
        return "";

    std::vector<char> str;
    str.resize(len);

    if (!file.read(str.data(), static_cast<int>(str.size())))
        return "";

    return std::string(str.data(), str.size());
}

std::string FooterVersion(uint32_t version)
{
    std::stringstream ss;
    if (version)
        ss << (version >> 4) << "." << (version & 0xf);
    return ss.str();
}

std::string FooterTime(uint64_t unix_time)
{
    auto t = static_cast<std::time_t>(unix_time);
    auto tm = *std::gmtime(&t);

    std::stringstream ss;
    if (unix_time)
        ss << std::put_time(&tm, "%Y-%m-%d %T");
    return ss.str();
}

class SCPDisk final : public DemandDisk
{
public:
    SCPDisk(bool normalised) : m_normalised(normalised) {}

    void add_track_data(const CylHead& cylhead, std::vector<std::vector<uint16_t>>&& trackdata)
    {
        m_data[cylhead] = std::move(trackdata);
        extend(cylhead);
    }

protected:
    TrackData load(const CylHead& cylhead, bool /*first_read*/) override
    {
        FluxData flux_revs;

        auto it = m_data.find(cylhead);
        if (it == m_data.end())
            return TrackData(cylhead);

        for (auto& rev_times : it->second)
        {
            std::vector<uint32_t> flux_times;
            flux_times.reserve(rev_times.size());

            uint32_t total_time = 0;
            for (auto time : rev_times) // note: big endian times!
            {
                if (!time)
                    total_time += 0x10000;
                else
                {
                    total_time += util::betoh(time);
                    flux_times.push_back(total_time * 25);  // 25ns sampling time
                    total_time = 0;
                }
            }

            flux_revs.push_back(std::move(flux_times));
        }

        return TrackData(cylhead, std::move(flux_revs), m_normalised);
    }

private:
    std::map<CylHead, std::vector<std::vector<uint16_t>>> m_data{};
    bool m_normalised = false;
};

bool ReadSCP(MemFile& file, std::shared_ptr<Disk>& disk)
{
    SCP_FILE_HEADER fh{};

    if (!file.rewind() || !file.read(&fh, sizeof(fh)) || std::string(fh.signature, 3) != "SCP")
        return false;

    if (!(fh.flags & FLAG_MODE) && fh.checksum)
    {
        auto checksum = std::accumulate(file.data().begin() + 0x10, file.data().end(), uint32_t(0));
        if (checksum != util::letoh(fh.checksum))
            Message(msgWarning, "file checksum is incorrect!");
    }

    /*if (!(fh.flags & FLAG_INDEX))
        throw util::exception("not an index-synchronised image");
    else*/ if (fh.revolutions == 0 || fh.revolutions > 10)
    throw util::exception("invalid revolution count (", fh.revolutions, ")");
    else if (fh.bitcell_width != 0 && fh.bitcell_width != 16)
    throw util::exception("unsupported bit cell width (", fh.bitcell_width, ")");
    else if (fh.start_track > fh.end_track)
    throw util::exception("start track (", fh.start_track, ") > end track (", fh.end_track, ")");

    std::vector<uint32_t> tdh_offsets(fh.end_track + 1);
    if (!file.read(tdh_offsets))
        throw util::exception("short file reading track offset index");

    auto scp_disk = std::make_shared<SCPDisk>((fh.flags & FLAG_TYPE) != 0);

    for (auto tracknr = fh.start_track; tracknr <= fh.end_track; ++tracknr)
    {
        TRACK_DATA_HEADER tdh;
        auto cyl = fh.heads == 0 ? (tracknr >> 1) : tracknr;
        auto head = fh.heads == 0 ? (tracknr & 1) : (fh.heads - 1);
        CylHead cylhead(cyl, head);

        if (!tdh_offsets[tracknr])
            continue;

        if (!file.seek(tdh_offsets[tracknr]) || !file.read(&tdh, sizeof(tdh)))
            throw util::exception("short file reading ", cylhead, " track header");
        else if (std::string(tdh.signature, 3) != "TRK")
            throw util::exception("invalid track signature on ", cylhead);
        else if (tdh.tracknr != tracknr)
            throw util::exception("track number mismatch (", tdh.tracknr, " != ", tracknr, ") in ", cylhead, " header");

        std::vector<uint32_t> rev_index(fh.revolutions * 3);
        if (!file.read(rev_index))
            throw util::exception("short file reading ", cylhead, " track index");

        std::vector<std::vector<uint16_t>> revs_data;
        revs_data.reserve(fh.revolutions);

        for (uint8_t rev = 0; rev < fh.revolutions; ++rev)
        {
            //          auto index_time  = util::letoh<uint32_t>(rev_index[rev * 3 + 0]);
            auto flux_count = util::letoh<uint32_t>(rev_index[rev * 3 + 1]);
            auto data_offset = util::letoh<uint32_t>(rev_index[rev * 3 + 2]);

            std::vector<uint16_t> flux_data(flux_count);
            if (!file.seek(tdh_offsets[tracknr] + data_offset) || !file.read(flux_data))
                throw util::exception("short error reading ", cylhead, " data");

            revs_data.push_back(std::move(flux_data));
        }

        scp_disk->add_track_data(cylhead, std::move(revs_data));
    }

    auto footer_offset = file.size() - static_cast<int>(sizeof(SCP_FILE_FOOTER));
    if ((fh.flags & FLAG_FOOTER) && footer_offset >= file.tell())
    {
        SCP_FILE_FOOTER ff{};

        if (file.seek(footer_offset) && file.read(&ff, sizeof(ff)) &&
            std::string(ff.sig, sizeof(ff.sig)) == "FPCS")
        {
            scp_disk->metadata["manufacturer"] = FooterString(file, ff.manufacturer_offset);
            scp_disk->metadata["model"] = FooterString(file, ff.model_offset);
            scp_disk->metadata["serial"] = FooterString(file, ff.serial_offset);
            scp_disk->metadata["creator"] = FooterString(file, ff.creator_offset);
            scp_disk->metadata["application"] = FooterString(file, ff.application_offset);
            scp_disk->metadata["comment"] = FooterString(file, ff.comments_offset);

            scp_disk->metadata["app_version"] = FooterVersion(ff.application_version);
            scp_disk->metadata["hw_version"] = FooterVersion(ff.hardware_version);
            scp_disk->metadata["fw_version"] = FooterVersion(ff.firmware_version);
            scp_disk->metadata["scp_version"] = FooterVersion(ff.format_revision);

            scp_disk->metadata["created"] = FooterTime(ff.creation_time);
            if (ff.modification_time != ff.creation_time)
                scp_disk->metadata["modified"] = FooterTime(ff.modification_time);
        }
    }
    else
    {
        scp_disk->metadata["app_version"] = FooterVersion(fh.revision);
        scp_disk->metadata["application"] = "SuperCard Pro software";

        std::stringstream ss;
        for (uint8_t b; file.read(&b, sizeof(b)) && std::isprint(b); )
            ss << static_cast<char>(b);
        scp_disk->metadata["created"] = ss.str();
    }

    scp_disk->metadata["index"] = (fh.flags & FLAG_INDEX) ? "synchronised" : "unsynchronised";
    scp_disk->metadata["tpi"] = (fh.flags & FLAG_TPI) ? "96 tpi" : "48 tpi";
    scp_disk->metadata["rpm"] = (fh.flags & FLAG_RPM) ? "360 rpm" : "300 rpm";
    scp_disk->metadata["quality"] = (fh.flags & FLAG_TYPE) ? "normalised" : "preservation";
    scp_disk->metadata["mode"] = (fh.flags & FLAG_MODE) ? "read/write" : "read-only";

    scp_disk->strType = "SCP";
    disk = scp_disk;

    return true;
}
<commit_msg>Changed SCP reading to scan all TDH entries<commit_after>// SuperCard Pro flux file format:
//  http://www.cbmstuff.com/downloads/scp/scp_image_specs.txt

#include "SAMdisk.h"
#include "DemandDisk.h"

constexpr auto DEFAULT_TDH_ENTRIES = 166;
constexpr auto MAX_TDH_ENTRIES = 168;

enum
{
    FLAG_INDEX = 1 << 0,    // set for index-synchronised, clear if not
    FLAG_TPI = 1 << 1,  // set for 96tpi, clear for 48tpi
    FLAG_RPM = 1 << 2,  // set for 360rpm, clear for 300rpm
    FLAG_TYPE = 1 << 3, // set for normalised flux, clear for preservation quality
    FLAG_MODE = 1 << 4, // set for read/write image, clear for read-only
    FLAG_FOOTER = 1 << 5    // set if footer block present, clear if absent
};

struct SCP_FILE_HEADER
{
    char signature[3];      // SCP
    uint8_t revision;       // (version << 4) | revision, 0x39 = v3.9
    uint8_t disk_type;
    uint8_t revolutions;    // number of stored revolutions
    uint8_t start_track;
    uint8_t end_track;
    uint8_t flags;
    uint8_t bitcell_width;  // 0 = 16 bits, non-zero = number of bits
    uint8_t heads;          // 0 = both heads, 1 = head 0 only, 2 = head 1 only
    uint8_t reserved;       // should be zero?
    uint32_t checksum;      // 32-bit checksum from after header to EOF (unless FLAG_MODE is set)
};

struct SCP_FILE_FOOTER
{
    uint32_t manufacturer_offset;
    uint32_t model_offset;
    uint32_t serial_offset;
    uint32_t creator_offset;
    uint32_t application_offset;
    uint32_t comments_offset;
    uint64_t creation_time;
    uint64_t modification_time;
    uint8_t application_version;
    uint8_t hardware_version;
    uint8_t firmware_version;
    uint8_t format_revision;
    char sig[4];
};

struct TRACK_DATA_HEADER
{
    char signature[3];      // TRK
    uint8_t tracknr;
};

std::string FooterString(MemFile& file, uint32_t offset)
{
    uint16_t len = 0;
    if (!offset || !file.seek(static_cast<int>(offset)) || !file.read(&len, sizeof(len)))
        return "";

    len = util::letoh(len);
    if (static_cast<int>(offset + len) >= file.size())
        return "";

    std::vector<char> str;
    str.resize(len);

    if (!file.read(str.data(), static_cast<int>(str.size())))
        return "";

    return std::string(str.data(), str.size());
}

std::string FooterVersion(uint32_t version)
{
    std::stringstream ss;
    if (version)
        ss << (version >> 4) << "." << (version & 0xf);
    return ss.str();
}

std::string FooterTime(uint64_t unix_time)
{
    auto t = static_cast<std::time_t>(unix_time);
    auto tm = *std::gmtime(&t);

    std::stringstream ss;
    if (unix_time)
        ss << std::put_time(&tm, "%Y-%m-%d %T");
    return ss.str();
}

class SCPDisk final : public DemandDisk
{
public:
    SCPDisk(bool normalised) : m_normalised(normalised) {}

    void add_track_data(const CylHead& cylhead, std::vector<std::vector<uint16_t>>&& trackdata)
    {
        m_data[cylhead] = std::move(trackdata);
        extend(cylhead);
    }

protected:
    TrackData load(const CylHead& cylhead, bool /*first_read*/) override
    {
        FluxData flux_revs;

        auto it = m_data.find(cylhead);
        if (it == m_data.end())
            return TrackData(cylhead);

        for (auto& rev_times : it->second)
        {
            std::vector<uint32_t> flux_times;
            flux_times.reserve(rev_times.size());

            uint32_t total_time = 0;
            for (auto time : rev_times) // note: big endian times!
            {
                if (!time)
                    total_time += 0x10000;
                else
                {
                    total_time += util::betoh(time);
                    flux_times.push_back(total_time * 25);  // 25ns sampling time
                    total_time = 0;
                }
            }

            flux_revs.push_back(std::move(flux_times));
        }

        return TrackData(cylhead, std::move(flux_revs), m_normalised);
    }

private:
    std::map<CylHead, std::vector<std::vector<uint16_t>>> m_data{};
    bool m_normalised = false;
};

bool ReadSCP(MemFile& file, std::shared_ptr<Disk>& disk)
{
    SCP_FILE_HEADER fh{};

    if (!file.rewind() || !file.read(&fh, sizeof(fh)) || std::string(fh.signature, 3) != "SCP")
        return false;

    if (!(fh.flags & FLAG_MODE) && fh.checksum)
    {
        auto checksum = std::accumulate(file.data().begin() + 0x10, file.data().end(), uint32_t(0));
        if (checksum != util::letoh(fh.checksum))
            Message(msgWarning, "file checksum is incorrect!");
    }

    /*if (!(fh.flags & FLAG_INDEX))
        throw util::exception("not an index-synchronised image");
    else*/ if (fh.revolutions == 0 || fh.revolutions > 10)
        throw util::exception("invalid revolution count (", fh.revolutions, ")");
    else if (fh.bitcell_width != 0 && fh.bitcell_width != 16)
        throw util::exception("unsupported bit cell width (", fh.bitcell_width, ")");

    auto entries = (fh.end_track == MAX_TDH_ENTRIES - 1) ? MAX_TDH_ENTRIES : DEFAULT_TDH_ENTRIES;
    std::vector<uint32_t> tdh_offsets(entries);
    if (!file.read(tdh_offsets))
        throw util::exception("short file reading track offset index");

    auto scp_disk = std::make_shared<SCPDisk>((fh.flags & FLAG_TYPE) != 0);

    for (int tracknr = 0; tracknr < static_cast<int>(tdh_offsets.size()); ++tracknr)
    {
        TRACK_DATA_HEADER tdh;
        CylHead cylhead(tracknr / 2, tracknr & 1);

        if (!tdh_offsets[tracknr])
            continue;

        if (!file.seek(tdh_offsets[tracknr]) || !file.read(&tdh, sizeof(tdh)))
            throw util::exception("short file reading ", cylhead, " track header");
        else if (std::string(tdh.signature, 3) != "TRK")
            throw util::exception("invalid track signature on ", cylhead);
        else if (tdh.tracknr != tracknr)
            throw util::exception("track number mismatch (", tdh.tracknr, " != ", tracknr, ") in ", cylhead, " header");

        std::vector<uint32_t> rev_index(fh.revolutions * 3);
        if (!file.read(rev_index))
            throw util::exception("short file reading ", cylhead, " track index");

        std::vector<std::vector<uint16_t>> revs_data;
        revs_data.reserve(fh.revolutions);

        for (uint8_t rev = 0; rev < fh.revolutions; ++rev)
        {
            // auto index_time  = util::letoh<uint32_t>(rev_index[rev * 3 + 0]);
            auto flux_count = util::letoh<uint32_t>(rev_index[rev * 3 + 1]);
            auto data_offset = util::letoh<uint32_t>(rev_index[rev * 3 + 2]);

            std::vector<uint16_t> flux_data(flux_count);
            if (!file.seek(tdh_offsets[tracknr] + data_offset) || !file.read(flux_data))
                throw util::exception("short error reading ", cylhead, " data");

            revs_data.push_back(std::move(flux_data));
        }

        scp_disk->add_track_data(cylhead, std::move(revs_data));
    }

    auto footer_offset = file.size() - static_cast<int>(sizeof(SCP_FILE_FOOTER));
    if ((fh.flags & FLAG_FOOTER) && footer_offset >= file.tell())
    {
        SCP_FILE_FOOTER ff{};

        if (file.seek(footer_offset) && file.read(&ff, sizeof(ff)) &&
            std::string(ff.sig, sizeof(ff.sig)) == "FPCS")
        {
            scp_disk->metadata["manufacturer"] = FooterString(file, ff.manufacturer_offset);
            scp_disk->metadata["model"] = FooterString(file, ff.model_offset);
            scp_disk->metadata["serial"] = FooterString(file, ff.serial_offset);
            scp_disk->metadata["creator"] = FooterString(file, ff.creator_offset);
            scp_disk->metadata["application"] = FooterString(file, ff.application_offset);
            scp_disk->metadata["comment"] = FooterString(file, ff.comments_offset);

            scp_disk->metadata["app_version"] = FooterVersion(ff.application_version);
            scp_disk->metadata["hw_version"] = FooterVersion(ff.hardware_version);
            scp_disk->metadata["fw_version"] = FooterVersion(ff.firmware_version);
            scp_disk->metadata["scp_version"] = FooterVersion(ff.format_revision);

            scp_disk->metadata["created"] = FooterTime(ff.creation_time);
            if (ff.modification_time != ff.creation_time)
                scp_disk->metadata["modified"] = FooterTime(ff.modification_time);
        }
    }
    else
    {
        scp_disk->metadata["app_version"] = FooterVersion(fh.revision);
        scp_disk->metadata["application"] = "SuperCard Pro software";

        std::stringstream ss;
        for (uint8_t b; file.read(&b, sizeof(b)) && std::isprint(b); )
            ss << static_cast<char>(b);
        scp_disk->metadata["created"] = ss.str();
    }

    scp_disk->metadata["index"] = (fh.flags & FLAG_INDEX) ? "synchronised" : "unsynchronised";
    scp_disk->metadata["tpi"] = (fh.flags & FLAG_TPI) ? "96 tpi" : "48 tpi";
    scp_disk->metadata["rpm"] = (fh.flags & FLAG_RPM) ? "360 rpm" : "300 rpm";
    scp_disk->metadata["quality"] = (fh.flags & FLAG_TYPE) ? "normalised" : "preservation";
    scp_disk->metadata["mode"] = (fh.flags & FLAG_MODE) ? "read/write" : "read-only";

    scp_disk->strType = "SCP";
    disk = scp_disk;

    return true;
}
<|endoftext|>


<commit_before>#include "StdAfx.h"

#define BUFFERSIZE (2 * 1024 * 1024)
char gBuffer[BUFFERSIZE];
int lastFrame = -1;

/*
    Construct Ultracomm object, connect to server,
    and set/check parameters.
*/
Ultracomm::Ultracomm(const UltracommOptions& myuopt)
    : uopt(myuopt),
      address(myuopt.opt["address"].as<string>()),
      acqmode(myuopt.opt["acqmode"].as<string>()),
      datatype(myuopt.opt["datatype"].as<int>()),
      verbose(myuopt.opt["verbose"].as<int>())
{
    connect();
    if (verbose) {
        cerr << "Setting data to acquire to datatype " << datatype << ".\n";
    }
    ult.setDataToAcquire(datatype);
    set_int_imaging_params();
    check_int_imaging_params();
    const int compstat = uopt.opt["compression_status"].as<int>();
    if (verbose) {
        cerr << "Setting compression status to " << compstat << ".\n";
    }
    if (acqmode == "continuous") {
        ult.setCallback(frame_callback);
    }
    if (! ult.setCompressionStatus(compstat)) {
        cerr << "Failed to set compression status to " << compstat << ".\n";
        throw ParameterMismatchError();
    }
}

/*
   Connect to the Ultrasonix.
*/
void Ultracomm::connect()
{
    if (verbose) {
        cerr << "Connecting to ultrasonix at address " << address << ".\n";
    }
    if (!ult.connect(address.c_str()))
    {
        throw ConnectionError();
    }
/*
TODO: throw different errors depending on problem. Note that FD_CONNECT error
is when server address is bad (e.g. 123), and SOCKET_ERROR is when we couldn't connect
(address good but not in research mode). Should also test when address is good
but ultrasound machine not running):

C:\build\ultracomm\bin\Debug>ultracomm --address 123 --output asdf
monitorConnection(): FD_CONNECT had ERROR
Could not connect to Ultrasonix.

C:\build\ultracomm\bin\Debug>ultracomm --address 192.168.1.200 --output asdf
monitorConnection(): WSAEnumNetworkEvents() ret SOCKET_ERROR (10093)
CmonitorConnection(): CONNECTION_COMM lost

ould not connect to Ultrasonix.
*/
}

/*
    Disconnect from Ultrasonix.
*/
void Ultracomm::disconnect()
{
    if (ult.isConnected())
    {
        if (verbose) {
            cerr << "Disconnecting from Ultrasonix.\n";
        }
        ult.disconnect();
    }
    else
    {
        if (verbose) {
            cerr << "Already disconnected from Ultrasonix.\n";
        }
    }
}

/*
  Put Ultrasonix into freeze state and wait for confirmation.
*/
void Ultracomm::wait_for_freeze()
{
    // 1 = FROZEN; 0 = IMAGING
    if (ult.getFreezeState() != 1)
    {
        if (verbose) {
            cerr << "Freezing Ultrasonix.\n";
        }
        ult.toggleFreeze();
    }
    else
    {
        if (verbose) {
            cerr << "Ultrasonix already frozen.\n";
        }
    }
    // Wait for server to acknowledge it has frozen.
    // TODO: this would be safer with a timeout.
    while (ult.getFreezeState() != 1)
    {
        if (verbose) {
            cerr << "Waiting for confirmation that Ultrasonix has frozen.\n";
        }
    }
    /*
       FIXME: Occasionally we get the message

           sendAndWait(): WaitForSingleObject() ret WAIT_TIMEOUT

       when retrieving cine data. When --verbose is turned on this message
       usually appears in save_data() after the call to getCineData() and
       before write(). According to the discussion at
       http://research.ultrasonix.com/viewtopic.php?f=5&t=1100&p=4245&hilit=delay#p4245
       adding a delay after toggleFreeze() can prevent this condition, so
       we add it here.

       Possibly this delay is not necessary when used with newer versions of
       the Ultrasonix library.
    */
    if (
           uopt.opt["ms_delay_after_freeze"].as<int>() > 0 &&
           !uopt.opt.count("freeze-only")
    )
    {
        Sleep(uopt.opt["ms_delay_after_freeze"].as<int>());
    }
}

/*
  Put ultrasonix into imaging state.
*/
void Ultracomm::wait_for_unfreeze()
{
    // 1 = FROZEN; 0 = IMAGING
    if (ult.getFreezeState() != 0)
    {
        if (verbose) {
            cerr << "Unfreezing Ultrasonix.\n";
        }
        ult.toggleFreeze();
    }
    else
    {
        if (verbose) {
            cerr << "Ultrasonix already imaging.\n";
        }
    }
    // Wait for server to acknowledge it has switched to imaging.
    // TODO: this would be safer with a timeout.
    while (ult.getFreezeState() != 0)
    {
        if (verbose) {
            cerr << "Waiting for confirmation that Ultrasonix is imaging.\n";
        }
    }
}

/*
    Print all ultrasonix parameters.
*/
void Ultracomm::dump_params()
{
    int i = 0;
    int val;
    uParam param;
    while (ult.getParam(i++, param))
    {
        ult.getParamValue(param.name, val);
        printf("%d\t%s\t%s\t%d\n", i, param.id, param.name, val);
    }
}

/*
    Set all integer-type Ultrasonix imaging parameters, as specified on the
    command line or in the parameter file.
*/
void Ultracomm::set_int_imaging_params()
{
    po::variables_map params = uopt.opt;
    po::options_description iopts = uopt.int_imaging_params;
    for ( auto iter = iopts.options().begin(); iter != iopts.options().end(); ++iter)
    {
        /*
           Ultracomm program options do not contain spaces. Some of the
           parameters used by the ulterius setParamValue() call do contain
           spaces, and none appear to use underscore. The convention is that
           ultracomm program options containing underscore correspond to
           ulterius parameters that have the same name, but with the
           underscores replaced with spaces (' ').

           optname: the name as used in ultracomm program options (underscores)
           ultname: the name as used by ulterius setParamValue() (spaces)
        */
        string optname = (*iter)->long_name();
        string ultname = boost::replace_all_copy(optname, "_", " ");
        if (params.count(optname)) {
            int val = params[optname].as<int>();
            if (verbose) {
                cerr << "Setting '" << ultname << "' to value " << val << ".\n";
            }
            ult.setParamValue(ultname.c_str(), val);
        }
    }
}

/*
    Verify that integer-type Ultrasonix imaging parameters have value as
    specified by user.
*/
void Ultracomm::check_int_imaging_params()
{
    po::variables_map params = uopt.opt;
    po::options_description iopts = uopt.int_imaging_params;
    for ( auto iter = iopts.options().begin(); iter != iopts.options().end(); ++iter)
    {
        /*
           See comments in set_int_imaging_params() for mangling of
           parameter names.
        */
        string optname = (*iter)->long_name();
        string ultname = boost::replace_all_copy(optname, "_", " ");
        if (params.count(optname)) {
            int expected, got;
            expected = params[optname].as<int>();
            if (verbose) {
                cerr << "Getting value of '" << ultname << "'. Expecting " << expected << ".\n";
            }
            ult.getParamValue(ultname.c_str(), got);
            if (got != expected) {
                cerr << "Parameter '" << ultname << "' expected " << expected << " and got " << got << ".\n";
                throw ParameterMismatchError();
            }
        }
    }
}

/*
    Get data from Ultrasonix and save to file.
*/
void Ultracomm::save_data()
{
    po::variables_map params = uopt.opt;
    const string outname = params["output"].as<string>();
    int num_frames = ult.getCineDataCount((uData)datatype);
    uDataDesc desc;
    if (! ult.getDataDescriptor((uData)datatype, desc))
    {
        throw DataDescriptorError();
    }
    if (! ult.isDataAvailable((uData)datatype))
    {
        throw DataError();
    }
    ofstream outfile (outname, ios::out | ios::binary);
    write_header(outfile, desc, num_frames);

    // TODO: figure out why buffer and sz makes program crash
    //const int sz = 2 * 1024 * 1024;
    //char buffer[BUFFERSIZE];  // TODO: determine appropriate sizes on the fly
//    int num_frames = ult.getCineDataCount((uData)datatype);

    // TODO: framesize assumes desc.ss is always a multiple of 8, and that might not be safe.
    int framesize = (desc.ss / 8) * desc.w * desc.h;
    for (int idx = 0; idx < num_frames; idx++)
    {
        if (verbose) {
            cerr << "Getting cine data for frame " << idx << ".\n";
        }
        ult.getCineData((uData)datatype, idx, false, (char**)&gBuffer, BUFFERSIZE);
        outfile.write(gBuffer, framesize);
        if (verbose) {
            cerr << "Wrote cine data for frame " << idx << ".\n";
        }
    }
    outfile.close();
}

/*
    TODO: header information is probably correct for .bpr files but possibly
    not for other datatypes.
*/
void Ultracomm::write_header(ofstream& outfile, const uDataDesc& desc, const int& num_frames)
{
    if (verbose) {
        cerr << "Writing header.\n";
    }
    // Integer fields that we can write directly. Luckily these all are
    // consecutive fields in the header.
    const int isize = sizeof(__int32);
    static const int fields[] = {
        datatype,
        num_frames,
        desc.w,
        desc.h,
        desc.ss,
        desc.roi.ulx,
        desc.roi.uly,
        desc.roi.urx,
        desc.roi.ury,
        desc.roi.brx,
        desc.roi.bry,
        desc.roi.blx,
        desc.roi.bly,
        uopt.opt["probe-id"].as<int>()
    };
    for (int i = 0; i < sizeof(fields) / sizeof(fields[0]); ++i) {
        outfile.write(reinterpret_cast<const char *>(&(__int32)fields[i]), isize);
    }

    // Integer fields that we have to query from Ultrasonix. These are also
    // consecutive in the header.
    // FIXME: Determine if these are the correct params to put in the header.
    static const string queries[] = {
        "b-freq",
        "vec-freq",
        "frame rate",
        "b-ldensity"
    };
    for (int i = 0; i < sizeof(queries) / sizeof(queries[0]); ++i) {
        int val;
        ult.getParamValue(queries[i].c_str(), val);
        outfile.write(reinterpret_cast<const char *>(&(__int32)val), isize);
    }
    // FIXME: Figure out how to determine the value of the 'extra' field.
    // It will probably be added to queries[], but I don't know the param name.
    // For now we'll hard code it with value 0.
    int extra = 0;
    outfile.write(reinterpret_cast<const char *>(&(__int32)extra), isize);
}

/*
    Callback.
*/
bool Ultracomm::frame_callback(void* data, int type, int sz, bool cine, int frmnum)
{
/*
    if (verbose) {
        cerr << "In callback.\n";
    }

    if (!data || !sz)
    {
// TODO: proper error
        printf("Error: no actual frame data received\n");
        return false;
    }

    if (BUFFERSIZE < sz)
    {
// TODO: proper error
        printf("Error: frame too large for current buffer\n");
        return false;
    }
*/

    //printf("[Rx] type:(%d) size:(%d) cine:(%d) gBuffer:(%d) frame:(%d)\n", type, sz, cine, &gBuffer, frmnum);
    if (frmnum != lastFrame + 1) {
        printf("Skipped frame(s) %d - %d.\n", lastFrame + 1, frmnum - 1);
    }
    lastFrame = frmnum;
    // make sure we dont do an operation that takes longer than the acquisition frame rate
    memcpy(gBuffer, data, sz);

    return true;
}
<commit_msg>Fix dump_params().<commit_after>#include "StdAfx.h"

#define BUFFERSIZE (2 * 1024 * 1024)
char gBuffer[BUFFERSIZE];
int lastFrame = -1;

/*
    Construct Ultracomm object, connect to server,
    and set/check parameters.
*/
Ultracomm::Ultracomm(const UltracommOptions& myuopt)
    : uopt(myuopt),
      address(myuopt.opt["address"].as<string>()),
      acqmode(myuopt.opt["acqmode"].as<string>()),
      datatype(myuopt.opt["datatype"].as<int>()),
      verbose(myuopt.opt["verbose"].as<int>())
{
    connect();
    if (verbose) {
        cerr << "Setting data to acquire to datatype " << datatype << ".\n";
    }
    ult.setDataToAcquire(datatype);
    set_int_imaging_params();
    check_int_imaging_params();
    const int compstat = uopt.opt["compression_status"].as<int>();
    if (verbose) {
        cerr << "Setting compression status to " << compstat << ".\n";
    }
    if (acqmode == "continuous") {
        ult.setCallback(frame_callback);
    }
    if (! ult.setCompressionStatus(compstat)) {
        cerr << "Failed to set compression status to " << compstat << ".\n";
        throw ParameterMismatchError();
    }
}

/*
   Connect to the Ultrasonix.
*/
void Ultracomm::connect()
{
    if (verbose) {
        cerr << "Connecting to ultrasonix at address " << address << ".\n";
    }
    if (!ult.connect(address.c_str()))
    {
        throw ConnectionError();
    }
/*
TODO: throw different errors depending on problem. Note that FD_CONNECT error
is when server address is bad (e.g. 123), and SOCKET_ERROR is when we couldn't connect
(address good but not in research mode). Should also test when address is good
but ultrasound machine not running):

C:\build\ultracomm\bin\Debug>ultracomm --address 123 --output asdf
monitorConnection(): FD_CONNECT had ERROR
Could not connect to Ultrasonix.

C:\build\ultracomm\bin\Debug>ultracomm --address 192.168.1.200 --output asdf
monitorConnection(): WSAEnumNetworkEvents() ret SOCKET_ERROR (10093)
CmonitorConnection(): CONNECTION_COMM lost

ould not connect to Ultrasonix.
*/
}

/*
    Disconnect from Ultrasonix.
*/
void Ultracomm::disconnect()
{
    if (ult.isConnected())
    {
        if (verbose) {
            cerr << "Disconnecting from Ultrasonix.\n";
        }
        ult.disconnect();
    }
    else
    {
        if (verbose) {
            cerr << "Already disconnected from Ultrasonix.\n";
        }
    }
}

/*
  Put Ultrasonix into freeze state and wait for confirmation.
*/
void Ultracomm::wait_for_freeze()
{
    // 1 = FROZEN; 0 = IMAGING
    if (ult.getFreezeState() != 1)
    {
        if (verbose) {
            cerr << "Freezing Ultrasonix.\n";
        }
        ult.toggleFreeze();
    }
    else
    {
        if (verbose) {
            cerr << "Ultrasonix already frozen.\n";
        }
    }
    // Wait for server to acknowledge it has frozen.
    // TODO: this would be safer with a timeout.
    while (ult.getFreezeState() != 1)
    {
        if (verbose) {
            cerr << "Waiting for confirmation that Ultrasonix has frozen.\n";
        }
    }
    /*
       FIXME: Occasionally we get the message

           sendAndWait(): WaitForSingleObject() ret WAIT_TIMEOUT

       when retrieving cine data. When --verbose is turned on this message
       usually appears in save_data() after the call to getCineData() and
       before write(). According to the discussion at
       http://research.ultrasonix.com/viewtopic.php?f=5&t=1100&p=4245&hilit=delay#p4245
       adding a delay after toggleFreeze() can prevent this condition, so
       we add it here.

       Possibly this delay is not necessary when used with newer versions of
       the Ultrasonix library.
    */
    if (
           uopt.opt["ms_delay_after_freeze"].as<int>() > 0 &&
           !uopt.opt.count("freeze-only")
    )
    {
        Sleep(uopt.opt["ms_delay_after_freeze"].as<int>());
    }
}

/*
  Put ultrasonix into imaging state.
*/
void Ultracomm::wait_for_unfreeze()
{
    // 1 = FROZEN; 0 = IMAGING
    if (ult.getFreezeState() != 0)
    {
        if (verbose) {
            cerr << "Unfreezing Ultrasonix.\n";
        }
        ult.toggleFreeze();
    }
    else
    {
        if (verbose) {
            cerr << "Ultrasonix already imaging.\n";
        }
    }
    // Wait for server to acknowledge it has switched to imaging.
    // TODO: this would be safer with a timeout.
    while (ult.getFreezeState() != 0)
    {
        if (verbose) {
            cerr << "Waiting for confirmation that Ultrasonix is imaging.\n";
        }
    }
}

/*
    Print all ultrasonix parameters.
*/
void Ultracomm::dump_params()
{
    int i = 0;
    int val;
    uParam param;
    cout << "index\tparam.id\tvalue\tparam.name\tparam.source\tparam.type\n";
    while (ult.getParam(i++, param))
    {
        // TODO: Learn about different param types and how to get their values.
        // This probably doesn't work properly for non-integer param values.
        // Have not been able to get getParamValue() to return false either.
        if (ult.getParamValue(param.id, val))
        {
            cout << i << "\t" << param.id << "\t" << val << "\t" << param.name << "\t" << param.source << "\t" << param.type << "\t" << param.unit << "\n";
        }
        else
        {
            cerr << i << "\t'" << param.id << "'\t" << param.name << "\tFAILED\n";
        }
    }
}

/*
    Set all integer-type Ultrasonix imaging parameters, as specified on the
    command line or in the parameter file.
*/
void Ultracomm::set_int_imaging_params()
{
    po::variables_map params = uopt.opt;
    po::options_description iopts = uopt.int_imaging_params;
    for ( auto iter = iopts.options().begin(); iter != iopts.options().end(); ++iter)
    {
        /*
           Ultracomm program options do not contain spaces. Some of the
           parameters used by the ulterius setParamValue() call do contain
           spaces, and none appear to use underscore. The convention is that
           ultracomm program options containing underscore correspond to
           ulterius parameters that have the same name, but with the
           underscores replaced with spaces (' ').

           optname: the name as used in ultracomm program options (underscores)
           ultname: the name as used by ulterius setParamValue() (spaces)
        */
        string optname = (*iter)->long_name();
        string ultname = boost::replace_all_copy(optname, "_", " ");
        if (params.count(optname)) {
            int val = params[optname].as<int>();
            if (verbose) {
                cerr << "Setting '" << ultname << "' to value " << val << ".\n";
            }
            ult.setParamValue(ultname.c_str(), val);
        }
    }
}

/*
    Verify that integer-type Ultrasonix imaging parameters have value as
    specified by user.
*/
void Ultracomm::check_int_imaging_params()
{
    po::variables_map params = uopt.opt;
    po::options_description iopts = uopt.int_imaging_params;
    for ( auto iter = iopts.options().begin(); iter != iopts.options().end(); ++iter)
    {
        /*
           See comments in set_int_imaging_params() for mangling of
           parameter names.
        */
        string optname = (*iter)->long_name();
        string ultname = boost::replace_all_copy(optname, "_", " ");
        if (params.count(optname)) {
            int expected, got;
            expected = params[optname].as<int>();
            ult.getParamValue(ultname.c_str(), got);
            if (verbose) {
                cerr << "Got value of '" << ultname << "'. Expected " << expected << " and got " << got << ".\n";
            }
            if (got != expected) {
                cerr << "Parameter '" << ultname << "' expected " << expected << " and got " << got << ".\n";
                throw ParameterMismatchError();
            }
        }
    }
}

/*
    Get data from Ultrasonix and save to file.
*/
void Ultracomm::save_data()
{
    po::variables_map params = uopt.opt;
    const string outname = params["output"].as<string>();
    int num_frames = ult.getCineDataCount((uData)datatype);
    uDataDesc desc;
    if (! ult.getDataDescriptor((uData)datatype, desc))
    {
        throw DataDescriptorError();
    }
    if (! ult.isDataAvailable((uData)datatype))
    {
        throw DataError();
    }
    ofstream outfile (outname, ios::out | ios::binary);
    write_header(outfile, desc, num_frames);

    // TODO: figure out why buffer and sz makes program crash
    //const int sz = 2 * 1024 * 1024;
    //char buffer[BUFFERSIZE];  // TODO: determine appropriate sizes on the fly
//    int num_frames = ult.getCineDataCount((uData)datatype);

    // TODO: framesize assumes desc.ss is always a multiple of 8, and that might not be safe.
    int framesize = (desc.ss / 8) * desc.w * desc.h;
    for (int idx = 0; idx < num_frames; idx++)
    {
        if (verbose) {
            cerr << "Getting cine data for frame " << idx << ".\n";
        }
        ult.getCineData((uData)datatype, idx, false, (char**)&gBuffer, BUFFERSIZE);
        outfile.write(gBuffer, framesize);
        if (verbose) {
            cerr << "Wrote cine data for frame " << idx << ".\n";
        }
    }
    outfile.close();
}

/*
    TODO: header information is probably correct for .bpr files but possibly
    not for other datatypes.
*/
void Ultracomm::write_header(ofstream& outfile, const uDataDesc& desc, const int& num_frames)
{
    if (verbose) {
        cerr << "Writing header.\n";
    }
    // Integer fields that we can write directly. Luckily these all are
    // consecutive fields in the header.
    const int isize = sizeof(__int32);
    static const int fields[] = {
        datatype,
        num_frames,
        desc.w,
        desc.h,
        desc.ss,
        desc.roi.ulx,
        desc.roi.uly,
        desc.roi.urx,
        desc.roi.ury,
        desc.roi.brx,
        desc.roi.bry,
        desc.roi.blx,
        desc.roi.bly,
        uopt.opt["probe-id"].as<int>()
    };
    for (int i = 0; i < sizeof(fields) / sizeof(fields[0]); ++i) {
        outfile.write(reinterpret_cast<const char *>(&(__int32)fields[i]), isize);
    }

    // Integer fields that we have to query from Ultrasonix. These are also
    // consecutive in the header.
    // FIXME: Determine if these are the correct params to put in the header.
    static const string queries[] = {
        "b-freq",
        "vec-freq",
        "frame rate",
        "b-ldensity"
    };
    for (int i = 0; i < sizeof(queries) / sizeof(queries[0]); ++i) {
        int val;
        ult.getParamValue(queries[i].c_str(), val);
        outfile.write(reinterpret_cast<const char *>(&(__int32)val), isize);
    }
    // FIXME: Figure out how to determine the value of the 'extra' field.
    // It will probably be added to queries[], but I don't know the param name.
    // For now we'll hard code it with value 0.
    int extra = 0;
    outfile.write(reinterpret_cast<const char *>(&(__int32)extra), isize);
}

/*
    Callback.
*/
bool Ultracomm::frame_callback(void* data, int type, int sz, bool cine, int frmnum)
{
/*
    if (verbose) {
        cerr << "In callback.\n";
    }

    if (!data || !sz)
    {
// TODO: proper error
        printf("Error: no actual frame data received\n");
        return false;
    }

    if (BUFFERSIZE < sz)
    {
// TODO: proper error
        printf("Error: frame too large for current buffer\n");
        return false;
    }
*/

    //printf("[Rx] type:(%d) size:(%d) cine:(%d) gBuffer:(%d) frame:(%d)\n", type, sz, cine, &gBuffer, frmnum);
    if (frmnum != lastFrame + 1) {
        printf("Skipped frame(s) %d - %d.\n", lastFrame + 1, frmnum - 1);
    }
    lastFrame = frmnum;
    // make sure we dont do an operation that takes longer than the acquisition frame rate
    memcpy(gBuffer, data, sz);

    return true;
}
<|endoftext|>