<commit_before>/**************************************************************************
 * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
 *                                                                        *
 * Author: The ALICE Off-line Project.                                    *
 * Contributors are mentioned in the code where appropriate.              *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, provided that the above copyright notice appears in all   *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purpose. It is          *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

//-------------------------------------------------------------------------
//               Implementation of the cascade vertexer class
//          Reads V0s and tracks, writes out cascade vertices
//                     Fills the ESD with the cascades
//    Origin: Christian Kuhn, IReS, Strasbourg, christian.kuhn@ires.in2p3.fr
//-------------------------------------------------------------------------

//modified by R. Vernet 30/6/2006 : daughter label
//modified by R. Vernet  3/7/2006 : causality
//modified by I. Belikov 24/11/2006 : static setter for the default cuts

#include <TObjArray.h>
#include <TTree.h>

#include "AliESDEvent.h"
#include "AliESDv0.h"
#include "AliESDcascade.h"
#include "AliCascadeVertexer.h"
#include "AliESDtrack.h"
#include "AliESDVertex.h"

ClassImp(AliCascadeVertexer)

//A set of loose cuts
Double_t
  AliCascadeVertexer::fgChi2max=33.;    //maximal allowed chi2
Double_t
  AliCascadeVertexer::fgDV0min=0.05;    //min V0 impact parameter
Double_t
  AliCascadeVertexer::fgMassWin=0.01;   //"window" around the Lambda mass
Double_t
  AliCascadeVertexer::fgDBachMin=0.035; //min bachelor impact parameter
Double_t
  AliCascadeVertexer::fgDCAmax=0.10;    //max DCA between the V0 and the track
Double_t
  AliCascadeVertexer::fgCPAmin=0.9985;  //min cosine of the cascade pointing angle
Double_t
  AliCascadeVertexer::fgRmin=0.2;       //min radius of the fiducial volume
Double_t
  AliCascadeVertexer::fgRmax=100.;      //max radius of the fiducial volume

Int_t AliCascadeVertexer::V0sTracks2CascadeVertices(AliESDEvent *event) {
  //--------------------------------------------------------------------
  // This function reconstructs cascade vertices
  //      Adapted to the ESD by I.Belikov (Jouri.Belikov@cern.ch)
  //--------------------------------------------------------------------
   const AliESDVertex *vtxSPD=event->GetVertex();
   const AliESDVertex *vtxT3D=event->GetPrimaryVertex();

   Double_t xPrimaryVertex=999, yPrimaryVertex=999, zPrimaryVertex=999;
   if (vtxT3D->GetStatus()) {
     xPrimaryVertex=vtxT3D->GetXv();
     yPrimaryVertex=vtxT3D->GetYv();
     zPrimaryVertex=vtxT3D->GetZv();
   }
   else {
     xPrimaryVertex=vtxSPD->GetXv();
     yPrimaryVertex=vtxSPD->GetYv();
     zPrimaryVertex=vtxSPD->GetZv();
   }

   Double_t b=event->GetMagneticField();
   Int_t nV0=(Int_t)event->GetNumberOfV0s();

   //stores relevant V0s in an array
   TObjArray vtcs(nV0);
   Int_t i;
   for (i=0; i<nV0; i++) {
       AliESDv0 *v=event->GetV0(i);
       if (v->GetOnFlyStatus()) continue;
       if (v->GetD(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex)<fDV0min) continue;
       vtcs.AddLast(v);
   }
   nV0=vtcs.GetEntriesFast();

   // stores relevant tracks in another array
   Int_t nentr=(Int_t)event->GetNumberOfTracks();
   TArrayI trk(nentr); Int_t ntr=0;
   for (i=0; i<nentr; i++) {
       AliESDtrack *esdtr=event->GetTrack(i);
       ULong_t status=esdtr->GetStatus();

       if ((status&AliESDtrack::kITSrefit)==0)
	  if ((status&AliESDtrack::kTPCrefit)==0) continue;

       if (TMath::Abs(esdtr->GetD(xPrimaryVertex,yPrimaryVertex,b))<fDBachMin) continue;

       trk[ntr++]=i;
   }

   Double_t massLambda=1.11568;
   Int_t ncasc=0;

   // Looking for the cascades...

   for (i=0; i<nV0; i++) { //loop on V0s

      AliESDv0 *v=(AliESDv0*)vtcs.UncheckedAt(i);
      v->ChangeMassHypothesis(kLambda0); // the v0 must be Lambda
      if (TMath::Abs(v->GetEffMass()-massLambda)>fMassWin) continue;

      for (Int_t j=0; j<ntr; j++) {//loop on tracks
	 Int_t bidx=trk[j];
 	 //Bo:   if (bidx==v->GetNindex()) continue; //bachelor and v0's negative tracks must be different
         if (bidx==v->GetIndex(0)) continue; //Bo:  consistency 0 for neg
	 AliESDtrack *btrk=event->GetTrack(bidx);
         if (btrk->GetSign()>0) continue;  // bachelor's charge

    	 AliESDv0 v0(*v), *pv0=&v0;
         AliExternalTrackParam bt(*btrk), *pbt=&bt;

         Double_t dca=PropagateToDCA(pv0,pbt,b);
         if (dca > fDCAmax) continue;

         AliESDcascade cascade(*pv0,*pbt,bidx);//constucts a cascade candidate
	 //PH        if (cascade.GetChi2Xi() > fChi2max) continue;

	 Double_t x,y,z; cascade.GetXYZcascade(x,y,z); // Bo: bug correction
         Double_t r2=x*x + y*y;
         if (r2 > fRmax*fRmax) continue;   // condition on fiducial zone
         if (r2 < fRmin*fRmin) continue;

	 Double_t pxV0,pyV0,pzV0;
	 pv0->GetPxPyPz(pxV0,pyV0,pzV0);
	 if (x*pxV0+y*pyV0+z*pzV0 < 0) continue; //causality

         Double_t x1,y1,z1; pv0->GetXYZ(x1,y1,z1);
         if (r2 > (x1*x1+y1*y1)) continue;

  	 if (cascade.GetCascadeCosineOfPointingAngle(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex) <fCPAmin) continue; //condition on the cascade pointing angle

	 event->AddCascade(&cascade);
         ncasc++;
      } // end loop tracks
   } // end loop V0s

   // Looking for the anti-cascades...

   for (i=0; i<nV0; i++) { //loop on V0s
      AliESDv0 *v=(AliESDv0*)vtcs.UncheckedAt(i);
      v->ChangeMassHypothesis(kLambda0Bar); //the v0 must be anti-Lambda
      if (TMath::Abs(v->GetEffMass()-massLambda)>fMassWin) continue;

      for (Int_t j=0; j<ntr; j++) {//loop on tracks
	 Int_t bidx=trk[j];
 	 //Bo:   if (bidx==v->GetPindex()) continue; //bachelor and v0's positive tracks must be different
         if (bidx==v->GetIndex(1)) continue; //Bo:  consistency 1 for pos
	 AliESDtrack *btrk=event->GetTrack(bidx);
         if (btrk->GetSign()<0) continue;  // bachelor's charge

	 AliESDv0 v0(*v), *pv0=&v0;
         AliESDtrack bt(*btrk), *pbt=&bt;

         Double_t dca=PropagateToDCA(pv0,pbt,b);
         if (dca > fDCAmax) continue;

         AliESDcascade cascade(*pv0,*pbt,bidx); //constucts a cascade candidate
	 //PH         if (cascade.GetChi2Xi() > fChi2max) continue;

	 Double_t x,y,z; cascade.GetXYZcascade(x,y,z); // Bo: bug correction
         Double_t r2=x*x + y*y;
         if (r2 > fRmax*fRmax) continue;   // condition on fiducial zone
         if (r2 < fRmin*fRmin) continue;

	 Double_t pxV0,pyV0,pzV0;
	 pv0->GetPxPyPz(pxV0,pyV0,pzV0);
	 if (x*pxV0+y*pyV0+z*pzV0 < 0) continue; //causality

         Double_t x1,y1,z1; pv0->GetXYZ(x1,y1,z1);
         if (r2 > (x1*x1+y1*y1)) continue;
         if (z*z > z1*z1) continue;

	 if (cascade.GetCascadeCosineOfPointingAngle(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex) < fCPAmin) continue; //condition on the cascade pointing angle
	 event->AddCascade(&cascade);
         ncasc++;

      } // end loop tracks
   } // end loop V0s

Info("V0sTracks2CascadeVertices","Number of reconstructed cascades: %d",ncasc);

   return 0;
}

Double_t det(Double_t a00, Double_t a01, Double_t a10, Double_t a11){
  // determinant 2x2
  return a00*a11 - a01*a10;
}

Double_t det (Double_t a00,Double_t a01,Double_t a02,
                     Double_t a10,Double_t a11,Double_t a12,
                     Double_t a20,Double_t a21,Double_t a22) {
  // determinant 3x3
  return
  a00*det(a11,a12,a21,a22)-a01*det(a10,a12,a20,a22)+a02*det(a10,a11,a20,a21);
}

Double_t AliCascadeVertexer::
PropagateToDCA(AliESDv0 *v, AliExternalTrackParam *t, Double_t b) {
  //--------------------------------------------------------------------
  // This function returns the DCA between the V0 and the track
  //--------------------------------------------------------------------
  Double_t alpha=t->GetAlpha(), cs1=TMath::Cos(alpha), sn1=TMath::Sin(alpha);
  Double_t r[3]; t->GetXYZ(r);
  Double_t x1=r[0], y1=r[1], z1=r[2];
  Double_t p[3]; t->GetPxPyPz(p);
  Double_t px1=p[0], py1=p[1], pz1=p[2];

  Double_t x2,y2,z2;     // position and momentum of V0
  Double_t px2,py2,pz2;

  v->GetXYZ(x2,y2,z2);
  v->GetPxPyPz(px2,py2,pz2);

// calculation dca

  Double_t dd= det(x2-x1,y2-y1,z2-z1,px1,py1,pz1,px2,py2,pz2);
  Double_t ax= det(py1,pz1,py2,pz2);
  Double_t ay=-det(px1,pz1,px2,pz2);
  Double_t az= det(px1,py1,px2,py2);

  Double_t dca=TMath::Abs(dd)/TMath::Sqrt(ax*ax + ay*ay + az*az);

//points of the DCA
  Double_t t1 = det(x2-x1,y2-y1,z2-z1,px2,py2,pz2,ax,ay,az)/
                det(px1,py1,pz1,px2,py2,pz2,ax,ay,az);

  x1 += px1*t1; y1 += py1*t1; //z1 += pz1*t1;

  //propagate track to the points of DCA

  x1=x1*cs1 + y1*sn1;
  if (!t->PropagateTo(x1,b)) {
    Error("PropagateToDCA","Propagation failed !");
    return 1.e+33;
  }

  return dca;
}

<commit_msg>Removing the unnecessary cut, and setting explicitly the DCA between the cascade daughters (A. Maire)<commit_after>/**************************************************************************
 * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
 *                                                                        *
 * Author: The ALICE Off-line Project.                                    *
 * Contributors are mentioned in the code where appropriate.              *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, provided that the above copyright notice appears in all   *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purpose. It is          *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

//-------------------------------------------------------------------------
//               Implementation of the cascade vertexer class
//          Reads V0s and tracks, writes out cascade vertices
//                     Fills the ESD with the cascades
//    Origin: Christian Kuhn, IReS, Strasbourg, christian.kuhn@ires.in2p3.fr
//-------------------------------------------------------------------------

//modified by R. Vernet 30/6/2006 : daughter label
//modified by R. Vernet  3/7/2006 : causality
//modified by I. Belikov 24/11/2006 : static setter for the default cuts

#include <TObjArray.h>
#include <TTree.h>

#include "AliESDEvent.h"
#include "AliESDv0.h"
#include "AliESDcascade.h"
#include "AliCascadeVertexer.h"
#include "AliESDtrack.h"
#include "AliESDVertex.h"

ClassImp(AliCascadeVertexer)

//A set of loose cuts
Double_t
  AliCascadeVertexer::fgChi2max=33.;    //maximal allowed chi2
Double_t
  AliCascadeVertexer::fgDV0min=0.05;    //min V0 impact parameter
Double_t
  AliCascadeVertexer::fgMassWin=0.01;   //"window" around the Lambda mass
Double_t
  AliCascadeVertexer::fgDBachMin=0.035; //min bachelor impact parameter
Double_t
  AliCascadeVertexer::fgDCAmax=0.10;    //max DCA between the V0 and the track
Double_t
  AliCascadeVertexer::fgCPAmin=0.9985;  //min cosine of the cascade pointing angle
Double_t
  AliCascadeVertexer::fgRmin=0.2;       //min radius of the fiducial volume
Double_t
  AliCascadeVertexer::fgRmax=100.;      //max radius of the fiducial volume

Int_t AliCascadeVertexer::V0sTracks2CascadeVertices(AliESDEvent *event) {
  //--------------------------------------------------------------------
  // This function reconstructs cascade vertices
  //      Adapted to the ESD by I.Belikov (Jouri.Belikov@cern.ch)
  //--------------------------------------------------------------------
   const AliESDVertex *vtxSPD=event->GetVertex();
   const AliESDVertex *vtxT3D=event->GetPrimaryVertex();

   Double_t xPrimaryVertex=999, yPrimaryVertex=999, zPrimaryVertex=999;
   if (vtxT3D->GetStatus()) {
     xPrimaryVertex=vtxT3D->GetXv();
     yPrimaryVertex=vtxT3D->GetYv();
     zPrimaryVertex=vtxT3D->GetZv();
   }
   else {
     xPrimaryVertex=vtxSPD->GetXv();
     yPrimaryVertex=vtxSPD->GetYv();
     zPrimaryVertex=vtxSPD->GetZv();
   }

   Double_t b=event->GetMagneticField();
   Int_t nV0=(Int_t)event->GetNumberOfV0s();

   //stores relevant V0s in an array
   TObjArray vtcs(nV0);
   Int_t i;
   for (i=0; i<nV0; i++) {
       AliESDv0 *v=event->GetV0(i);
       if (v->GetOnFlyStatus()) continue;
       if (v->GetD(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex)<fDV0min) continue;
       vtcs.AddLast(v);
   }
   nV0=vtcs.GetEntriesFast();

   // stores relevant tracks in another array
   Int_t nentr=(Int_t)event->GetNumberOfTracks();
   TArrayI trk(nentr); Int_t ntr=0;
   for (i=0; i<nentr; i++) {
       AliESDtrack *esdtr=event->GetTrack(i);
       ULong_t status=esdtr->GetStatus();

       if ((status&AliESDtrack::kITSrefit)==0)
	  if ((status&AliESDtrack::kTPCrefit)==0) continue;

       if (TMath::Abs(esdtr->GetD(xPrimaryVertex,yPrimaryVertex,b))<fDBachMin) continue;

       trk[ntr++]=i;
   }

   Double_t massLambda=1.11568;
   Int_t ncasc=0;

   // Looking for the cascades...

   for (i=0; i<nV0; i++) { //loop on V0s

      AliESDv0 *v=(AliESDv0*)vtcs.UncheckedAt(i);
      v->ChangeMassHypothesis(kLambda0); // the v0 must be Lambda
      if (TMath::Abs(v->GetEffMass()-massLambda)>fMassWin) continue;

      for (Int_t j=0; j<ntr; j++) {//loop on tracks
	 Int_t bidx=trk[j];
 	 //Bo:   if (bidx==v->GetNindex()) continue; //bachelor and v0's negative tracks must be different
         if (bidx==v->GetIndex(0)) continue; //Bo:  consistency 0 for neg
	 AliESDtrack *btrk=event->GetTrack(bidx);
         if (btrk->GetSign()>0) continue;  // bachelor's charge

    	 AliESDv0 v0(*v), *pv0=&v0;
         AliExternalTrackParam bt(*btrk), *pbt=&bt;

         Double_t dca=PropagateToDCA(pv0,pbt,b);
         if (dca > fDCAmax) continue;

         AliESDcascade cascade(*pv0,*pbt,bidx);//constucts a cascade candidate
	 //PH        if (cascade.GetChi2Xi() > fChi2max) continue;

	 Double_t x,y,z; cascade.GetXYZcascade(x,y,z); // Bo: bug correction
         Double_t r2=x*x + y*y;
         if (r2 > fRmax*fRmax) continue;   // condition on fiducial zone
         if (r2 < fRmin*fRmin) continue;

	 Double_t pxV0,pyV0,pzV0;
	 pv0->GetPxPyPz(pxV0,pyV0,pzV0);
	 if (x*pxV0+y*pyV0+z*pzV0 < 0) continue; //causality

         Double_t x1,y1,z1; pv0->GetXYZ(x1,y1,z1);
         if (r2 > (x1*x1+y1*y1)) continue;

  	 if (cascade.GetCascadeCosineOfPointingAngle(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex) <fCPAmin) continue; //condition on the cascade pointing angle

         cascade.SetDcaXiDaughters(dca);
	 event->AddCascade(&cascade);
         ncasc++;
      } // end loop tracks
   } // end loop V0s

   // Looking for the anti-cascades...

   for (i=0; i<nV0; i++) { //loop on V0s
      AliESDv0 *v=(AliESDv0*)vtcs.UncheckedAt(i);
      v->ChangeMassHypothesis(kLambda0Bar); //the v0 must be anti-Lambda
      if (TMath::Abs(v->GetEffMass()-massLambda)>fMassWin) continue;

      for (Int_t j=0; j<ntr; j++) {//loop on tracks
	 Int_t bidx=trk[j];
 	 //Bo:   if (bidx==v->GetPindex()) continue; //bachelor and v0's positive tracks must be different
         if (bidx==v->GetIndex(1)) continue; //Bo:  consistency 1 for pos
	 AliESDtrack *btrk=event->GetTrack(bidx);
         if (btrk->GetSign()<0) continue;  // bachelor's charge

	 AliESDv0 v0(*v), *pv0=&v0;
         AliESDtrack bt(*btrk), *pbt=&bt;

         Double_t dca=PropagateToDCA(pv0,pbt,b);
         if (dca > fDCAmax) continue;

         AliESDcascade cascade(*pv0,*pbt,bidx); //constucts a cascade candidate
	 //PH         if (cascade.GetChi2Xi() > fChi2max) continue;

	 Double_t x,y,z; cascade.GetXYZcascade(x,y,z); // Bo: bug correction
         Double_t r2=x*x + y*y;
         if (r2 > fRmax*fRmax) continue;   // condition on fiducial zone
         if (r2 < fRmin*fRmin) continue;

	 Double_t pxV0,pyV0,pzV0;
	 pv0->GetPxPyPz(pxV0,pyV0,pzV0);
	 if (x*pxV0+y*pyV0+z*pzV0 < 0) continue; //causality

         Double_t x1,y1,z1; pv0->GetXYZ(x1,y1,z1);
         if (r2 > (x1*x1+y1*y1)) continue;

	 if (cascade.GetCascadeCosineOfPointingAngle(xPrimaryVertex,yPrimaryVertex,zPrimaryVertex) < fCPAmin) continue; //condition on the cascade pointing angle

         cascade.SetDcaXiDaughters(dca);
	 event->AddCascade(&cascade);
         ncasc++;

      } // end loop tracks
   } // end loop V0s

Info("V0sTracks2CascadeVertices","Number of reconstructed cascades: %d",ncasc);

   return 0;
}

Double_t det(Double_t a00, Double_t a01, Double_t a10, Double_t a11){
  // determinant 2x2
  return a00*a11 - a01*a10;
}

Double_t det (Double_t a00,Double_t a01,Double_t a02,
                     Double_t a10,Double_t a11,Double_t a12,
                     Double_t a20,Double_t a21,Double_t a22) {
  // determinant 3x3
  return
  a00*det(a11,a12,a21,a22)-a01*det(a10,a12,a20,a22)+a02*det(a10,a11,a20,a21);
}

Double_t AliCascadeVertexer::
PropagateToDCA(AliESDv0 *v, AliExternalTrackParam *t, Double_t b) {
  //--------------------------------------------------------------------
  // This function returns the DCA between the V0 and the track
  //--------------------------------------------------------------------
  Double_t alpha=t->GetAlpha(), cs1=TMath::Cos(alpha), sn1=TMath::Sin(alpha);
  Double_t r[3]; t->GetXYZ(r);
  Double_t x1=r[0], y1=r[1], z1=r[2];
  Double_t p[3]; t->GetPxPyPz(p);
  Double_t px1=p[0], py1=p[1], pz1=p[2];

  Double_t x2,y2,z2;     // position and momentum of V0
  Double_t px2,py2,pz2;

  v->GetXYZ(x2,y2,z2);
  v->GetPxPyPz(px2,py2,pz2);

// calculation dca

  Double_t dd= det(x2-x1,y2-y1,z2-z1,px1,py1,pz1,px2,py2,pz2);
  Double_t ax= det(py1,pz1,py2,pz2);
  Double_t ay=-det(px1,pz1,px2,pz2);
  Double_t az= det(px1,py1,px2,py2);

  Double_t dca=TMath::Abs(dd)/TMath::Sqrt(ax*ax + ay*ay + az*az);

//points of the DCA
  Double_t t1 = det(x2-x1,y2-y1,z2-z1,px2,py2,pz2,ax,ay,az)/
                det(px1,py1,pz1,px2,py2,pz2,ax,ay,az);

  x1 += px1*t1; y1 += py1*t1; //z1 += pz1*t1;

  //propagate track to the points of DCA

  x1=x1*cs1 + y1*sn1;
  if (!t->PropagateTo(x1,b)) {
    Error("PropagateToDCA","Propagation failed !");
    return 1.e+33;
  }

  return dca;
}

<|endoftext|>


<commit_before>/*
 * Copyright 2009, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * 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.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "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 COPYRIGHT
 * OWNER 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.
 */

// This file implements the OpenGL versions of the render surface resources,
// as well as the related GAPIGL function implementations.

#include "command_buffer/service/cross/gl/gapi_gl.h"
#include "command_buffer/service/cross/gl/render_surface_gl.h"

namespace o3d {
namespace command_buffer {

RenderSurfaceGL::RenderSurfaceGL(int width,
                                 int height,
                                 int mip_level,
                                 int side,
                                 TextureGL *texture)
    : width_(width), height_(height), mip_level_(mip_level), texture_(texture) {
}

RenderSurfaceGL* RenderSurfaceGL::Create(int width,
                                         int height,
                                         int mip_level,
                                         int side,
                                         TextureGL *texture) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);
  DCHECK_GE(mip_level, 0);
  DCHECK(texture);

  RenderSurfaceGL* render_surface =
      new RenderSurfaceGL(width, height, mip_level, side, texture);
  return render_surface;
}

RenderDepthStencilSurfaceGL::RenderDepthStencilSurfaceGL(
    int width,
    int height)
    : width_(width), height_(height) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);

  // If packed depth stencil is supported, create only one buffer for both
  // depth and stencil.
  if (GLEW_EXT_packed_depth_stencil) {
    glGenRenderbuffersEXT(1, render_buffers_);
    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[0]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_DEPTH24_STENCIL8_EXT,
                             width,
                             height);
    CHECK_GL_ERROR();
    render_buffers_[1] = render_buffers_[0];
  } else {
    glGenRenderbuffersEXT(2, render_buffers_);
    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[0]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_DEPTH_COMPONENT24,
                             width,
                             height);
    CHECK_GL_ERROR();

    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[1]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_STENCIL_INDEX8_EXT,
                             width,
                             height);
    CHECK_GL_ERROR();
  }
}

RenderDepthStencilSurfaceGL* RenderDepthStencilSurfaceGL::Create(
    int width,
    int height) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);

  return new RenderDepthStencilSurfaceGL(height, width);
}

// Copies the data from a texture resource.
BufferSyncInterface::ParseError GAPIGL::CreateRenderSurface(
    ResourceID id,
    unsigned int width,
    unsigned int height,
    unsigned int mip_level,
    unsigned int side,
    ResourceID texture_id) {
  if (id == current_surface_id_) {
    // This will delete the current surface which would be bad.
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }
  TextureGL *texture = textures_.Get(texture_id);
  if (!texture->render_surfaces_enabled()) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  } else {
    RenderSurfaceGL* render_surface = RenderSurfaceGL::Create(width,
                                                              height,
                                                              mip_level,
                                                              side,
                                                              texture);
    if (render_surface == NULL) {
      return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
    }
    render_surfaces_.Assign(id, render_surface);
  }
  return BufferSyncInterface::PARSE_NO_ERROR;
}

BufferSyncInterface::ParseError GAPIGL::DestroyRenderSurface(ResourceID id) {
  if (id == current_surface_id_) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }
  return render_surfaces_.Destroy(id) ?
      BufferSyncInterface::PARSE_NO_ERROR :
      BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
}

BufferSyncInterface::ParseError GAPIGL::CreateDepthSurface(
    ResourceID id,
    unsigned int width,
    unsigned int height) {
  if (id == current_depth_surface_id_) {
    // This will delete the current surface which would be bad.
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }
  RenderDepthStencilSurfaceGL* depth_surface =
    RenderDepthStencilSurfaceGL::Create(width, height);
  if (depth_surface == NULL) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }
  depth_surfaces_.Assign(id, depth_surface);
  return BufferSyncInterface::PARSE_NO_ERROR;
}

BufferSyncInterface::ParseError GAPIGL::DestroyDepthSurface(ResourceID id) {
  if (id == current_depth_surface_id_) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }
  return depth_surfaces_.Destroy(id) ?
      BufferSyncInterface::PARSE_NO_ERROR :
      BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
}

void ResetBoundAttachments() {
#ifdef _DEBUG
    GLint bound_framebuffer;
    ::glGetIntegerv(GL_FRAMEBUFFER_BINDING_EXT, &bound_framebuffer);
    DCHECK(bound_framebuffer != 0);
#endif

  // Reset the bound attachments to the current framebuffer object.
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_COLOR_ATTACHMENT0_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_DEPTH_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_STENCIL_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  CHECK_GL_ERROR();
}

bool BindDepthStencilBuffer(const RenderDepthStencilSurfaceGL* gl_surface) {
  // Bind both the depth and stencil attachments.
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_DEPTH_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 gl_surface->depth_buffer());
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_STENCIL_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 gl_surface->stencil_buffer());

  // Check for errors.
  GLenum framebuffer_status = ::glCheckFramebufferStatusEXT(GL_FRAMEBUFFER_EXT);
  if (GL_FRAMEBUFFER_COMPLETE_EXT != framebuffer_status) {
    return false;
  }

  CHECK_GL_ERROR();
  return true;
}

BufferSyncInterface::ParseError GAPIGL::SetRenderSurface(
    ResourceID render_surface_id,
    ResourceID depth_stencil_id) {
  if (render_surfaces_.Get(render_surface_id) == NULL &&
      depth_surfaces_.Get(depth_stencil_id) == NULL) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }

  ::glBindFramebufferEXT(GL_FRAMEBUFFER, render_surface_framebuffer_);
  ResetBoundAttachments();

  RenderSurfaceGL* render_surface = render_surfaces_.Get(render_surface_id);
  RenderDepthStencilSurfaceGL* depth_surface =
      depth_surfaces_.Get(render_surface_id);

  if (!render_surface->texture()->
          InstallFrameBufferObjects(render_surface) ||
      !BindDepthStencilBuffer(depth_surface)) {
    return BufferSyncInterface::PARSE_INVALID_ARGUMENTS;
  }

  // RenderSurface rendering is performed with an inverted Y, so the front
  // face winding must be changed to clock-wise.  See comments for
  // UpdateHelperConstant.
  glFrontFace(GL_CW);

  current_surface_id_ = render_surface_id;
  current_depth_surface_id_ = depth_stencil_id;
  return BufferSyncInterface::PARSE_NO_ERROR;
}

void GAPIGL::SetBackSurfaces() {
  // Bind the default context, and restore the default front-face winding.
  ::glBindFramebufferEXT(GL_FRAMEBUFFER, 0);
  glFrontFace(GL_CCW);
}

}  // namespace command_buffer
}  // namespace o3d

<commit_msg>fix for gl cb<commit_after>/*
 * Copyright 2009, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * 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.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "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 COPYRIGHT
 * OWNER 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.
 */

// This file implements the OpenGL versions of the render surface resources,
// as well as the related GAPIGL function implementations.

#include "command_buffer/service/cross/gl/gapi_gl.h"
#include "command_buffer/service/cross/gl/render_surface_gl.h"

namespace o3d {
namespace command_buffer {

RenderSurfaceGL::RenderSurfaceGL(int width,
                                 int height,
                                 int mip_level,
                                 int side,
                                 TextureGL *texture)
    : width_(width), height_(height), mip_level_(mip_level), texture_(texture) {
}

RenderSurfaceGL* RenderSurfaceGL::Create(int width,
                                         int height,
                                         int mip_level,
                                         int side,
                                         TextureGL *texture) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);
  DCHECK_GE(mip_level, 0);
  DCHECK(texture);

  RenderSurfaceGL* render_surface =
      new RenderSurfaceGL(width, height, mip_level, side, texture);
  return render_surface;
}

RenderDepthStencilSurfaceGL::RenderDepthStencilSurfaceGL(
    int width,
    int height)
    : width_(width), height_(height) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);

  // If packed depth stencil is supported, create only one buffer for both
  // depth and stencil.
  if (GLEW_EXT_packed_depth_stencil) {
    glGenRenderbuffersEXT(1, render_buffers_);
    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[0]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_DEPTH24_STENCIL8_EXT,
                             width,
                             height);
    CHECK_GL_ERROR();
    render_buffers_[1] = render_buffers_[0];
  } else {
    glGenRenderbuffersEXT(2, render_buffers_);
    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[0]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_DEPTH_COMPONENT24,
                             width,
                             height);
    CHECK_GL_ERROR();

    glBindRenderbufferEXT(GL_RENDERBUFFER_EXT, render_buffers_[1]);
    glRenderbufferStorageEXT(GL_RENDERBUFFER_EXT,
                             GL_STENCIL_INDEX8_EXT,
                             width,
                             height);
    CHECK_GL_ERROR();
  }
}

RenderDepthStencilSurfaceGL* RenderDepthStencilSurfaceGL::Create(
    int width,
    int height) {
  DCHECK_GT(width, 0);
  DCHECK_GT(height, 0);

  return new RenderDepthStencilSurfaceGL(height, width);
}

// Copies the data from a texture resource.
BufferSyncInterface::ParseError GAPIGL::CreateRenderSurface(
    ResourceID id,
    unsigned int width,
    unsigned int height,
    unsigned int mip_level,
    unsigned int side,
    ResourceID texture_id) {
  if (id == current_surface_id_) {
    // This will delete the current surface which would be bad.
    return BufferSyncInterface::kParseInvalidArguments;
  }
  TextureGL *texture = textures_.Get(texture_id);
  if (!texture->render_surfaces_enabled()) {
    return BufferSyncInterface::kParseInvalidArguments;
  } else {
    RenderSurfaceGL* render_surface = RenderSurfaceGL::Create(width,
                                                              height,
                                                              mip_level,
                                                              side,
                                                              texture);
    if (render_surface == NULL) {
      return BufferSyncInterface::kParseInvalidArguments;
    }
    render_surfaces_.Assign(id, render_surface);
  }
  return BufferSyncInterface::kParseNoError;
}

BufferSyncInterface::ParseError GAPIGL::DestroyRenderSurface(ResourceID id) {
  if (id == current_surface_id_) {
    return BufferSyncInterface::kParseInvalidArguments;
  }
  return render_surfaces_.Destroy(id) ?
      BufferSyncInterface::kParseNoError :
      BufferSyncInterface::kParseInvalidArguments;
}

BufferSyncInterface::ParseError GAPIGL::CreateDepthSurface(
    ResourceID id,
    unsigned int width,
    unsigned int height) {
  if (id == current_depth_surface_id_) {
    // This will delete the current surface which would be bad.
    return BufferSyncInterface::kParseInvalidArguments;
  }
  RenderDepthStencilSurfaceGL* depth_surface =
    RenderDepthStencilSurfaceGL::Create(width, height);
  if (depth_surface == NULL) {
    return BufferSyncInterface::kParseInvalidArguments;
  }
  depth_surfaces_.Assign(id, depth_surface);
  return BufferSyncInterface::kParseNoError;
}

BufferSyncInterface::ParseError GAPIGL::DestroyDepthSurface(ResourceID id) {
  if (id == current_depth_surface_id_) {
    return BufferSyncInterface::kParseInvalidArguments;
  }
  return depth_surfaces_.Destroy(id) ?
      BufferSyncInterface::kParseNoError :
      BufferSyncInterface::kParseInvalidArguments;
}

void ResetBoundAttachments() {
#ifdef _DEBUG
    GLint bound_framebuffer;
    ::glGetIntegerv(GL_FRAMEBUFFER_BINDING_EXT, &bound_framebuffer);
    DCHECK(bound_framebuffer != 0);
#endif

  // Reset the bound attachments to the current framebuffer object.
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_COLOR_ATTACHMENT0_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_DEPTH_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_STENCIL_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 0);

  CHECK_GL_ERROR();
}

bool BindDepthStencilBuffer(const RenderDepthStencilSurfaceGL* gl_surface) {
  // Bind both the depth and stencil attachments.
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_DEPTH_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 gl_surface->depth_buffer());
  ::glFramebufferRenderbufferEXT(GL_FRAMEBUFFER_EXT,
                                 GL_STENCIL_ATTACHMENT_EXT,
                                 GL_RENDERBUFFER_EXT,
                                 gl_surface->stencil_buffer());

  // Check for errors.
  GLenum framebuffer_status = ::glCheckFramebufferStatusEXT(GL_FRAMEBUFFER_EXT);
  if (GL_FRAMEBUFFER_COMPLETE_EXT != framebuffer_status) {
    return false;
  }

  CHECK_GL_ERROR();
  return true;
}

BufferSyncInterface::ParseError GAPIGL::SetRenderSurface(
    ResourceID render_surface_id,
    ResourceID depth_stencil_id) {
  if (render_surfaces_.Get(render_surface_id) == NULL &&
      depth_surfaces_.Get(depth_stencil_id) == NULL) {
    return BufferSyncInterface::kParseInvalidArguments;
  }

  ::glBindFramebufferEXT(GL_FRAMEBUFFER, render_surface_framebuffer_);
  ResetBoundAttachments();

  RenderSurfaceGL* render_surface = render_surfaces_.Get(render_surface_id);
  RenderDepthStencilSurfaceGL* depth_surface =
      depth_surfaces_.Get(render_surface_id);

  if (!render_surface->texture()->
          InstallFrameBufferObjects(render_surface) ||
      !BindDepthStencilBuffer(depth_surface)) {
    return BufferSyncInterface::kParseInvalidArguments;
  }

  // RenderSurface rendering is performed with an inverted Y, so the front
  // face winding must be changed to clock-wise.  See comments for
  // UpdateHelperConstant.
  glFrontFace(GL_CW);

  current_surface_id_ = render_surface_id;
  current_depth_surface_id_ = depth_stencil_id;
  return BufferSyncInterface::kParseNoError;
}

void GAPIGL::SetBackSurfaces() {
  // Bind the default context, and restore the default front-face winding.
  ::glBindFramebufferEXT(GL_FRAMEBUFFER, 0);
  glFrontFace(GL_CCW);
}

}  // namespace command_buffer
}  // namespace o3d

<|endoftext|>


<commit_before>#include "libtg.hpp"

using namespace std;

// argv[0]
string exe;

void usage() {
  cerr << "Usage mode: " << exe << " [options...] \n";
  cerr << endl;
  cerr << "I/O options:\n";
  cerr << "\t-i\t<file path>\tRead formula from file.\n";
  cerr << "\t-oinfo\t<file path>\tAppend formula information to file.\n";
  cerr << "\t-ocnf\t<file path>\tWrite CNF to file.\n";
  cerr << "\t-oproof\t<file path>\tWrite proof formula to file.\n";
  cerr << endl;
  cerr << "Pipeline options:\n";
  cerr << "\t-f\tFlat ((p|(q|r)|s) ---> (p|q|r|s)). Before DAG.\n";
  cerr << "\t-m\tMinimize DAG edges ((p&q)|(p&r&q) ---> (p&q)|((p&q)&r)).\n";
  cerr << "\t-a\t<algorithm number>\tSelect renaming algorithm.\n";
  cerr << "\t-s\tSimplify CNF. (see note 1)\n";
  cerr << "\t-ionly\tOnly compute formula information.\n";
  cerr << endl;
  cerr << "Additional options:\n";
  cerr << "\t-h\tDisplay usage mode.\n";
  cerr << "\t-syntax\tDisplay formula syntax.\n";
  cerr << endl;
  cerr << "Renaming algorithms (see note 2):\n";
  cerr << "\t1\tKnapsack 0-1 based algorithm.\n";
  cerr << "\t2\tBoy de la Tour's algorithm.\n";
  cerr << endl;
  cerr << "Notes:\n";
  cerr << "\t1) CNF simplification removes tautologies, repeated literals\n";
  cerr << "\tand repeated clauses.\n";
  cerr << "\t2) If option -a is not set or an invalid algorithm is passed,\n";
  cerr << "\tthen no renaming will happen at all.\n";
  cerr << "\t3) Formula information format:\n";
  cerr << "\t<file path>,<size>,<number of clauses>,<number of symbols>,\n";
  cerr << endl;
  exit(-1);
}

void syntax() {
  cerr << exe << ": Formula syntax.\n";
  cerr << endl;
  cerr << "Propositional symbols are strings of one or more of the following\n";
  cerr << "ASCII characters:\n";
  cerr << "\t_\t(underscore)\n";
  cerr << "\t0 to 9\t(decimal digits)\n";
  cerr << "\ta to z\t(lowercase letters)\n";
  cerr << "\tA to Z\t(uppercase letters)\n";
  cerr << endl;
  cerr << "Logic connectives are the following strings of ASCII characters:\n";
  cerr << "\t~\t(tilde)\t\t\t\t\t\tNegation.\n";
  cerr << "\t&\t(ampersand)\t\t\t\t\tConjunction.\n";
  cerr << "\t|\t(pipe)\t\t\t\t\t\tDisjunction.\n";
  cerr << "\t=>\t(equals and greater than signs)\t\t\tImplication.\n";
  cerr << "\t<=>\t(less than, equals and greater than signs)\tEquivalence.\n";
  cerr << endl;
  cerr << "Notes:\n";
  cerr << "\t1) Only one line is read by the parser.\n";
  cerr << "\t2) Whitespaces and multiple parentheses are completely ignored.\n";
  cerr << "\tExample:\n";
  cerr << "\t\t(((p => q)))\tis the same as\t\tp=>q\n";
  cerr << endl;
  exit(-1);
}

class Args {
  private:
    int argc;
    char** argv;
    unordered_map<string,int> idx;
  public:
    void init(int argc, char** argv) {
      this->argc = argc;
      this->argv = argv;
      for (int i = argc-1; 1 <= i; i--) idx[argv[i]] = i;
    }
    bool find(const string& str) const {
      return idx.find(str) != idx.end();
    }
    template <typename T>
    T opt(const string& name) const {
      auto it = idx.find(name);
      if (it == idx.end()) usage();
      int i = it->second+1;
      if (argc <= i) usage();
      stringstream ss;
      ss << argv[i];
      T ret;
      ss >> ret;
      if (ss.fail()) usage();
      return ret;
    }
    void checkflag(const string& name) const {
      if (!find(name)) usage();
    }
    template <typename T>
    void checkopt(const string& name, bool required = false) const {
      if (find(name)) opt<T>(name);
      else if (required) usage();
    }
} args;

// files
unordered_map<string,fstream*> files;
ostream& get_output_stream(const string& name, bool append = false) {
  if (!args.find("-o"+name)) return cout;
  string fn = args.opt<string>("-o"+name);
  auto it = files.find(fn);
  if (it != files.end()) return (ostream&)*it->second;
  fstream::openmode md = fstream::out;
  if (append) md = md | fstream::app;
  fstream* file = new fstream(fn,md);
  if (!file->is_open()) {
    cerr << exe << ": ";
    cerr << "error opening `" << fn << "' for " << name << " output.\n";
    exit(-1);
  }
  files[fn] = file;
  return (ostream&)*file;
}
void close_files() {
  for (auto& kv : files) {
    kv.second->close();
    delete kv.second;
  }
}

int main(int argc, char** argv) {
  // init args
  exe = argv[0];
  args.init(argc,argv);

  // display help
  if (args.find("-h")) usage();
  if (args.find("-syntax")) syntax();

  // init input stream
  string in_fn = "stdin";
  if (args.find("-i")) {
    in_fn = args.opt<string>("-i");
    if (!freopen(in_fn.c_str(),"r",stdin)) {
      cerr << exe << ": `" << in_fn << "' not found.\n";
      return -1;
    }
    reverse(in_fn.begin(),in_fn.end());
    auto pos = in_fn.find("/");
    if (pos != string::npos) in_fn = in_fn.substr(0,pos);
    reverse(in_fn.begin(),in_fn.end());
  }

  // init output streams
  ostream& info_stream  = get_output_stream("info", true);
  ostream& cnf_stream   = get_output_stream("cnf");
  ostream& proof_stream = get_output_stream("proof");

  // additional argument checks
  args.checkopt<int>("-a");

  // input
  getline(cin,raw);
  parse();
  stringstream ss;
  ss << T;
  string original = ss.str();

  // info only
  if (args.find("-ionly")) {
    info_stream << in_fn << ",";
    info_stream << size(T) << ",";
    info_stream << clauses(T) << ",";
    info_stream << symbols(T) << ",\n";
    cnf_stream << original << endl;
    close_files();
    return 0;
  }

  // pipeline
  nnf();
  if (args.find("-f")) flat();
  dag();
  if (args.find("-m")) mindag();
  if (args.find("-a")) switch (args.opt<int>("-a")) {
    case 1: knapsack();     break;
    case 2: boydelatour();  break;
  }
  rename();
  if (args.find("-s")) simplecnf();
  else cnf();

  // output
  info_stream << in_fn << ",";
  info_stream << CNF.size() << ",";
  info_stream << CNF.clauses() << ",";
  info_stream << CNF.symbols() << ",\n";
  cnf_stream << CNF << endl;
  proof_stream << "(" << original << ") <-> (" << CNF << ")\n";

  close_files();

  return 0;
}
<commit_msg>Adapting output<commit_after>#include "libtg.hpp"

using namespace std;

// argv[0]
string exe;

void usage() {
  cerr << "Usage mode: " << exe << " [options...] \n";
  cerr << endl;
  cerr << "I/O options:\n";
  cerr << "\t-i\t<file path>\tRead formula from file.\n";
  cerr << "\t-oinfo\t<file path>\tAppend formula information to file.\n";
  cerr << "\t-ocnf\t<file path>\tWrite CNF to file.\n";
  cerr << "\t-oproof\t<file path>\tWrite proof formula to file.\n";
  cerr << endl;
  cerr << "Pipeline options:\n";
  cerr << "\t-f\tFlat ((p|(q|r)|s) ---> (p|q|r|s)). Before DAG.\n";
  cerr << "\t-m\tMinimize DAG edges ((p&q)|(p&r&q) ---> (p&q)|((p&q)&r)).\n";
  cerr << "\t-a\t<algorithm number>\tSelect renaming algorithm.\n";
  cerr << "\t-s\tSimplify CNF. (see note 1)\n";
  cerr << "\t-ionly\tOnly compute formula information.\n";
  cerr << endl;
  cerr << "Additional options:\n";
  cerr << "\t-h\tDisplay usage mode.\n";
  cerr << "\t-syntax\tDisplay formula syntax.\n";
  cerr << endl;
  cerr << "Renaming algorithms (see note 2):\n";
  cerr << "\t1\tKnapsack 0-1 based algorithm.\n";
  cerr << "\t2\tBoy de la Tour's algorithm.\n";
  cerr << endl;
  cerr << "Notes:\n";
  cerr << "\t1) CNF simplification removes tautologies, repeated literals\n";
  cerr << "\tand repeated clauses.\n";
  cerr << "\t2) If option -a is not set or an invalid algorithm is passed,\n";
  cerr << "\tthen no renaming will happen at all.\n";
  cerr << "\t3) Formula information format:\n";
  cerr << "\t<file path>,<size>,<number of clauses>,<number of symbols>,\n";
  cerr << endl;
  exit(-1);
}

void syntax() {
  cerr << exe << ": Formula syntax.\n";
  cerr << endl;
  cerr << "Propositional symbols are strings of one or more of the following\n";
  cerr << "ASCII characters:\n";
  cerr << "\t_\t(underscore)\n";
  cerr << "\t0 to 9\t(decimal digits)\n";
  cerr << "\ta to z\t(lowercase letters)\n";
  cerr << "\tA to Z\t(uppercase letters)\n";
  cerr << endl;
  cerr << "Logic connectives are the following strings of ASCII characters:\n";
  cerr << "\t~\t(tilde)\t\t\t\t\t\tNegation.\n";
  cerr << "\t&\t(ampersand)\t\t\t\t\tConjunction.\n";
  cerr << "\t|\t(pipe)\t\t\t\t\t\tDisjunction.\n";
  cerr << "\t=>\t(equals and greater than signs)\t\t\tImplication.\n";
  cerr << "\t<=>\t(less than, equals and greater than signs)\tEquivalence.\n";
  cerr << endl;
  cerr << "Notes:\n";
  cerr << "\t1) Only one line is read by the parser.\n";
  cerr << "\t2) Whitespaces and multiple parentheses are completely ignored.\n";
  cerr << "\tExample:\n";
  cerr << "\t\t(((p => q)))\tis the same as\t\tp=>q\n";
  cerr << endl;
  exit(-1);
}

class Args {
  private:
    int argc;
    char** argv;
    unordered_map<string,int> idx;
  public:
    void init(int argc, char** argv) {
      this->argc = argc;
      this->argv = argv;
      for (int i = argc-1; 1 <= i; i--) idx[argv[i]] = i;
    }
    bool find(const string& str) const {
      return idx.find(str) != idx.end();
    }
    template <typename T>
    T opt(const string& name) const {
      auto it = idx.find(name);
      if (it == idx.end()) usage();
      int i = it->second+1;
      if (argc <= i) usage();
      stringstream ss;
      ss << argv[i];
      T ret;
      ss >> ret;
      if (ss.fail()) usage();
      return ret;
    }
    void checkflag(const string& name) const {
      if (!find(name)) usage();
    }
    template <typename T>
    void checkopt(const string& name, bool required = false) const {
      if (find(name)) opt<T>(name);
      else if (required) usage();
    }
} args;

// files
unordered_map<string,fstream*> files;
ostream& get_output_stream(const string& name, bool append = false) {
  if (!args.find("-o"+name)) return cout;
  string fn = args.opt<string>("-o"+name);
  auto it = files.find(fn);
  if (it != files.end()) return (ostream&)*it->second;
  fstream::openmode md = fstream::out;
  if (append) md = md | fstream::app;
  fstream* file = new fstream(fn,md);
  if (!file->is_open()) {
    cerr << exe << ": ";
    cerr << "error opening `" << fn << "' for " << name << " output.\n";
    exit(-1);
  }
  files[fn] = file;
  return (ostream&)*file;
}
void close_files() {
  for (auto& kv : files) {
    kv.second->close();
    delete kv.second;
  }
}

int main(int argc, char** argv) {
  // init args
  exe = argv[0];
  args.init(argc,argv);

  // display help
  if (args.find("-h")) usage();
  if (args.find("-syntax")) syntax();

  // init input stream
  string in_fn = "stdin";
  if (args.find("-i")) {
    in_fn = args.opt<string>("-i");
    if (!freopen(in_fn.c_str(),"r",stdin)) {
      cerr << exe << ": `" << in_fn << "' not found.\n";
      return -1;
    }
    reverse(in_fn.begin(),in_fn.end());
    auto pos = in_fn.find("/");
    if (pos != string::npos) in_fn = in_fn.substr(0,pos);
    reverse(in_fn.begin(),in_fn.end());
  }

  // init output streams
  ostream& info_stream  = get_output_stream("info", true);
  ostream& cnf_stream   = get_output_stream("cnf");
  ostream& proof_stream = get_output_stream("proof");

  // additional argument checks
  args.checkopt<int>("-a");

  // input
  getline(cin,raw);
  parse();
  stringstream ss;
  ss << T;
  string original = ss.str();

  // info only
  if (args.find("-ionly")) {
    info_stream << in_fn << ",";
    info_stream << size(T) << ",";
    info_stream << clauses(T) << ",";
    info_stream << symbols(T) << ",\n";
    cnf_stream << original << endl;
    close_files();
    return 0;
  }

  // pipeline
  nnf();
  if (args.find("-f")) flat();
  dag();
  if (args.find("-m")) mindag();
  if (args.find("-a")) switch (args.opt<int>("-a")) {
    case 1: knapsack();     break;
    case 2: boydelatour();  break;
  }
  rename();
  if (args.find("-s")) simplecnf();
  else cnf();

  // output
  info_stream << in_fn << ",";
  info_stream << CNF.size() << ",";
  info_stream << CNF.clauses() << ",";
  info_stream << CNF.symbols() << ",\n";
  cnf_stream << CNF << endl;
  proof_stream << "-((" << original << ") <-> (" << CNF << "))\n";

  close_files();

  return 0;
}
<|endoftext|>


<commit_before>/*
 *  Copyright (c) 2014-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */

#define _WIN32_DCOM
#define WIN32_LEAN_AND_MEAN

#include <Windows.h>
/// clang-format off
#include <LM.h>
#include <sddl.h>
// clang-format on

#include <iterator>
#include <map>
#include <string>

#include <boost/algorithm/hex.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>
#include <boost/range/adaptor/map.hpp>
#include <boost/range/algorithm.hpp>

#include <osquery/core.h>
#include <osquery/filesystem.h>
#include <osquery/logger.h>
#include <osquery/tables.h>

#include "osquery/core/conversions.h"
#include "osquery/core/windows/wmi.h"
#include "osquery/filesystem/fileops.h"
#include "osquery/tables/system/windows/registry.h"

namespace fs = boost::filesystem;

namespace osquery {
namespace tables {

const std::set<int> kRegistryStringTypes = {
    REG_SZ, REG_MULTI_SZ, REG_EXPAND_SZ};

const std::map<std::string, HKEY> kRegistryHives = {
    {"HKEY_CLASSES_ROOT", HKEY_CLASSES_ROOT},
    {"HKEY_CURRENT_CONFIG", HKEY_CURRENT_CONFIG},
    {"HKEY_CURRENT_USER", HKEY_CURRENT_USER},
    {"HKEY_CURRENT_USER_LOCAL_SETTINGS", HKEY_CURRENT_USER_LOCAL_SETTINGS},
    {"HKEY_LOCAL_MACHINE", HKEY_LOCAL_MACHINE},
    {"HKEY_PERFORMANCE_DATA", HKEY_PERFORMANCE_DATA},
    {"HKEY_PERFORMANCE_NLSTEXT", HKEY_PERFORMANCE_NLSTEXT},
    {"HKEY_PERFORMANCE_TEXT", HKEY_PERFORMANCE_TEXT},
    {"HKEY_USERS", HKEY_USERS},
};

const std::map<DWORD, std::string> kRegistryTypes = {
    {REG_BINARY, "REG_BINARY"},
    {REG_DWORD, "REG_DWORD"},
    {REG_DWORD_BIG_ENDIAN, "REG_DWORD_BIG_ENDIAN"},
    {REG_EXPAND_SZ, "REG_EXPAND_SZ"},
    {REG_LINK, "REG_LINK"},
    {REG_MULTI_SZ, "REG_MULTI_SZ"},
    {REG_NONE, "REG_NONE"},
    {REG_QWORD, "REG_QWORD"},
    {REG_SZ, "REG_SZ"},
    {REG_FULL_RESOURCE_DESCRIPTOR, "REG_FULL_RESOURCE_DESCRIPTOR"},
    {REG_RESOURCE_LIST, "REG_RESOURCE_LIST"},
};

Status getUsernameFromKey(const std::string& key, std::string& rUsername) {
  if (!boost::starts_with(key, "HKEY_USERS")) {
    return Status(1, "Can not extract username from non-HKEY_USERS key");
  }

  auto toks = osquery::split(key, kRegSep);
  if (toks.size() < 2) {
    return Status(
        1, "Improperly-formatted HKEY_USERS key, cannot extract username");
  }

  PSID sid;
  if (!ConvertStringSidToSidA(toks[1].c_str(), &sid)) {
    return Status(GetLastError(), "Could not convert string to sid");
  } else {
    wchar_t accntName[UNLEN] = {0};
    wchar_t domName[DNLEN] = {0};
    unsigned long accntNameLen = UNLEN;
    unsigned long domNameLen = DNLEN;
    SID_NAME_USE eUse;
    if (!LookupAccountSidW(nullptr,
                           sid,
                           accntName,
                           &accntNameLen,
                           domName,
                           &domNameLen,
                           &eUse)) {
      return Status(GetLastError(), "Could not find sid");
    } else {
      rUsername = std::move(wstringToString(accntName));
    }
  }
  return Status(0, "OK");
}

inline void explodeRegistryPath(const std::string& path,
                                std::string& rHive,
                                std::string& rKey) {
  auto toks = osquery::split(path, kRegSep);
  rHive = toks.front();
  toks.erase(toks.begin());
  rKey = osquery::join(toks, kRegSep);
}

/// Microsoft helper function for getting the contents of a registry key
void queryKey(const std::string& keyPath, QueryData& results) {
  std::string hive;
  std::string key;
  explodeRegistryPath(keyPath, hive, key);

  if (kRegistryHives.count(hive) != 1) {
    return;
  }

  HKEY hRegistryHandle;
  auto ret = RegOpenKeyEx(kRegistryHives.at(hive),
                          TEXT(key.c_str()),
                          0,
                          KEY_READ,
                          &hRegistryHandle);

  if (ret != ERROR_SUCCESS) {
    return;
  }

  const DWORD maxKeyLength = 255;
  const DWORD maxValueName = 16383;
  TCHAR achClass[MAX_PATH] = TEXT("");
  DWORD cchClassName = MAX_PATH;
  DWORD cSubKeys = 0;
  DWORD cbMaxSubKey;
  DWORD cchMaxClass;
  DWORD cValues;
  DWORD cchMaxValueName;
  DWORD cbMaxValueData;
  DWORD cbSecurityDescriptor;
  DWORD retCode;
  FILETIME ftLastWriteTime;
  retCode = RegQueryInfoKey(hRegistryHandle,
                            achClass,
                            &cchClassName,
                            nullptr,
                            &cSubKeys,
                            &cbMaxSubKey,
                            &cchMaxClass,
                            &cValues,
                            &cchMaxValueName,
                            &cbMaxValueData,
                            &cbSecurityDescriptor,
                            &ftLastWriteTime);

  TCHAR achKey[maxKeyLength];
  DWORD cbName;

  // Process registry subkeys
  if (cSubKeys > 0) {
    for (DWORD i = 0; i < cSubKeys; i++) {
      cbName = maxKeyLength;
      retCode = RegEnumKeyEx(hRegistryHandle,
                             i,
                             achKey,
                             &cbName,
                             nullptr,
                             nullptr,
                             nullptr,
                             &ftLastWriteTime);
      if (retCode != ERROR_SUCCESS) {
        continue;
      }
      Row r;
      r["key"] = keyPath;
      r["type"] = "subkey";
      r["name"] = achKey;
      r["path"] = keyPath + kRegSep + achKey;
      r["mtime"] = std::to_string(osquery::filetimeToUnixtime(ftLastWriteTime));
      results.push_back(r);
    }
  }

  if (cValues <= 0) {
    return;
  }

  DWORD cchValue = maxKeyLength;
  TCHAR achValue[maxValueName];
  BYTE* bpDataBuff =
      (cbMaxValueData == 0) ? nullptr : new BYTE[cbMaxValueData]();

  // Process registry values
  for (size_t i = 0, retCode = ERROR_SUCCESS; i < cValues; i++) {
    size_t cnt = 0;
    cchValue = maxValueName;
    achValue[0] = '\0';

    retCode = RegEnumValue(hRegistryHandle,
                           static_cast<DWORD>(i),
                           achValue,
                           &cchValue,
                           nullptr,
                           nullptr,
                           nullptr,
                           nullptr);
    if (retCode != ERROR_SUCCESS) {
      continue;
    }

    DWORD lpData = cbMaxValueData;
    DWORD lpType;

    retCode = RegQueryValueEx(
        hRegistryHandle, achValue, 0, &lpType, bpDataBuff, &lpData);
    if (retCode != ERROR_SUCCESS) {
      continue;
    }

    // It's possible for registry entries to have been inserted incorrectly
    // resulting in non-null-terminated strings
    if (bpDataBuff != nullptr && lpData != 0 &&
        kRegistryStringTypes.find(lpType) != kRegistryStringTypes.end()) {
      bpDataBuff[lpData - 1] = 0x00;
    }

    Row r;
    r["key"] = keyPath;
    r["name"] = ((achValue[0] == '\0') ? "(Default)" : achValue);
    r["path"] = keyPath + kRegSep + achValue;
    if (kRegistryTypes.count(lpType) > 0) {
      r["type"] = kRegistryTypes.at(lpType);
    } else {
      r["type"] = "UNKNOWN";
    }
    r["mtime"] = std::to_string(osquery::filetimeToUnixtime(ftLastWriteTime));

    if (bpDataBuff != nullptr) {
      /// REG_LINK is a Unicode string, which in Windows is wchar_t
      char* regLinkStr = nullptr;
      if (lpType == REG_LINK) {
        regLinkStr = new char[cbMaxValueData];
        const size_t newSize = cbMaxValueData;
        size_t convertedChars = 0;
        wcstombs_s(&convertedChars,
                   regLinkStr,
                   newSize,
                   (wchar_t*)bpDataBuff,
                   _TRUNCATE);
      }

      BYTE* bpDataBuffTmp = bpDataBuff;
      std::vector<std::string> multiSzStrs;
      std::vector<char> regBinary;
      std::string data;

      switch (lpType) {
      case REG_FULL_RESOURCE_DESCRIPTOR:
      case REG_RESOURCE_LIST:
      case REG_BINARY:
        for (size_t i = 0; i < cbMaxValueData; i++) {
          regBinary.push_back((char)bpDataBuff[i]);
        }
        boost::algorithm::hex(
            regBinary.begin(), regBinary.end(), std::back_inserter(data));
        r["data"] = data;
        break;
      case REG_DWORD:
        r["data"] = std::to_string(*((int*)bpDataBuff));
        break;
      case REG_DWORD_BIG_ENDIAN:
        r["data"] = std::to_string(_byteswap_ulong(*((int*)bpDataBuff)));
        break;
      case REG_EXPAND_SZ:
        r["data"] = std::string((char*)bpDataBuff);
        break;
      case REG_LINK:
        r["data"] = std::string(regLinkStr);
        break;
      case REG_MULTI_SZ:
        while (*bpDataBuffTmp != 0x00) {
          std::string s((char*)bpDataBuffTmp);
          bpDataBuffTmp += s.size() + 1;
          multiSzStrs.push_back(s);
        }
        r["data"] = boost::algorithm::join(multiSzStrs, ",");
        break;
      case REG_NONE:
        r["data"] = "(zero-length binary value)";
        break;
      case REG_QWORD:
        r["data"] = std::to_string(*((unsigned long long*)bpDataBuff));
        break;
      case REG_SZ:
        r["data"] = std::string((char*)bpDataBuff);
        break;
      default:
        r["data"] = "";
        break;
      }
      if (regLinkStr != nullptr) {
        delete[](regLinkStr);
      }
      ZeroMemory(bpDataBuff, cbMaxValueData);
    }
    results.push_back(r);
  }
  if (bpDataBuff != nullptr) {
    delete[](bpDataBuff);
  }
  RegCloseKey(hRegistryHandle);
}

static inline void populateDefaultKeys(std::set<std::string>& rKeys) {
  boost::copy(kRegistryHives | boost::adaptors::map_keys,
              std::inserter(rKeys, rKeys.end()));
}

static inline void populateSubkeys(std::set<std::string>& rKeys,
                                   bool replaceKeys = false) {
  std::set<std::string> newKeys;
  if (!replaceKeys) {
    newKeys = rKeys;
  }

  for (const auto& key : rKeys) {
    QueryData regResults;
    queryKey(key, regResults);
    for (const auto& r : regResults) {
      if (r.at("type") == "subkey") {
        newKeys.insert(r.at("path"));
      }
    }
  }
  rKeys = std::move(newKeys);
}

static inline void appendSubkeyToKeys(const std::string& subkey,
                                      std::set<std::string>& rKeys) {
  std::set<std::string> newKeys{};
  for (auto& key : rKeys) {
    newKeys.insert(std::move(key) + kRegSep + subkey);
  }
  rKeys = std::move(newKeys);
}

static inline Status populateAllKeysRecursive(
    std::set<std::string>& rKeys,
    size_t currDepth = 1,
    size_t maxDepth = kRegMaxRecursiveDepth) {
  if (currDepth > maxDepth) {
    return Status(1, "Max recursive depth reached");
  }

  auto size_pre = rKeys.size();
  populateSubkeys(rKeys);
  if (size_pre < rKeys.size()) {
    auto status = populateAllKeysRecursive(rKeys, ++currDepth);
    if (!status.ok()) {
      return status;
    }
  }

  return Status(0, "OK");
}

Status expandRegistryGlobs(const std::string& pattern,
                           std::set<std::string>& results) {
  auto pathElems = osquery::split(pattern, kRegSep);
  if (pathElems.size() == 0) {
    return Status(0, "OK");
  }

  /*
   * Pattern is '%%', grab everything.
   * Note that if '%%' is present but not at the end of the pattern,
   * then it is treated like a single glob.
   */
  if (boost::ends_with(pathElems[0], kSQLGlobRecursive) &&
      pathElems.size() == 1) {
    populateDefaultKeys(results);
    return populateAllKeysRecursive(results);
  }

  // Special handling to insert default keys when glob present in first elem
  if (pathElems[0].find(kSQLGlobWildcard) != std::string::npos) {
    populateDefaultKeys(results);
    pathElems.erase(pathElems.begin());
  } else {
    results.insert(pathElems[0]);
    pathElems.erase(pathElems.begin());
  }

  for (auto& elem = pathElems.begin(); elem != pathElems.end(); ++elem) {
    // We only care about  a recursive glob if it comes at the end of the
    // pattern i.e. 'HKEY_LOCAL_MACHINE\SOFTWARE\%%'
    if (boost::ends_with(*elem, kSQLGlobRecursive) &&
        *elem == pathElems.back()) {
      return populateAllKeysRecursive(results);
    } else if ((*elem).find(kSQLGlobWildcard) != std::string::npos) {
      populateSubkeys(results, true);
    } else {
      appendSubkeyToKeys(*elem, results);
    }
  }
  return Status(0, "OK");
}

static inline void maybeWarnLocalUsers(const std::set<std::string>& rKeys) {
  std::string hive, _;
  for (const auto& key : rKeys) {
    explodeRegistryPath(key, hive, _);
    if (hive == "HKEY_CURRENT_USER" ||
        hive == "HKEY_CURRENT_USER_LOCAL_SETTINGS") {
      LOG(WARNING) << "CURRENT_USER hives are not queryable by osqueryd; "
                      "query HKEY_USERS with the desired users SID instead";
      break;
    }
  }
}

QueryData genRegistry(QueryContext& context) {
  QueryData results;
  std::set<std::string> keys;

  if (!(context.hasConstraint("key", EQUALS) ||
        context.hasConstraint("key", LIKE))) {
    // We default to display all HIVEs
    expandRegistryGlobs(kSQLGlobWildcard, keys);
  } else {
    keys = context.constraints["key"].getAll(EQUALS);
    auto status = context.expandConstraints(
        "key",
        LIKE,
        keys,
        ([&](const std::string& pattern, std::set<std::string>& out) {
          std::set<std::string> resolvedKeys;
          auto status = expandRegistryGlobs(pattern, resolvedKeys);
          out.insert(resolvedKeys.begin(), resolvedKeys.end());
          return status;
        }));
    if (!status.ok()) {
      LOG(INFO) << "Failed to expand globs: " + status.getMessage();
    }
  }

  maybeWarnLocalUsers(keys);

  for (const auto& key : keys) {
    queryKey(key, results);
  }
  return results;
}
}
}
<commit_msg>[Fix 3249] Addressing permissive build fail (#3250)<commit_after>/*
 *  Copyright (c) 2014-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */

#define _WIN32_DCOM
#define WIN32_LEAN_AND_MEAN

#include <Windows.h>
/// clang-format off
#include <LM.h>
#include <sddl.h>
// clang-format on

#include <iterator>
#include <map>
#include <string>

#include <boost/algorithm/hex.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>
#include <boost/range/adaptor/map.hpp>
#include <boost/range/algorithm.hpp>

#include <osquery/core.h>
#include <osquery/filesystem.h>
#include <osquery/logger.h>
#include <osquery/tables.h>

#include "osquery/core/conversions.h"
#include "osquery/core/windows/wmi.h"
#include "osquery/filesystem/fileops.h"
#include "osquery/tables/system/windows/registry.h"

namespace fs = boost::filesystem;

namespace osquery {
namespace tables {

const std::set<int> kRegistryStringTypes = {
    REG_SZ, REG_MULTI_SZ, REG_EXPAND_SZ};

const std::map<std::string, HKEY> kRegistryHives = {
    {"HKEY_CLASSES_ROOT", HKEY_CLASSES_ROOT},
    {"HKEY_CURRENT_CONFIG", HKEY_CURRENT_CONFIG},
    {"HKEY_CURRENT_USER", HKEY_CURRENT_USER},
    {"HKEY_CURRENT_USER_LOCAL_SETTINGS", HKEY_CURRENT_USER_LOCAL_SETTINGS},
    {"HKEY_LOCAL_MACHINE", HKEY_LOCAL_MACHINE},
    {"HKEY_PERFORMANCE_DATA", HKEY_PERFORMANCE_DATA},
    {"HKEY_PERFORMANCE_NLSTEXT", HKEY_PERFORMANCE_NLSTEXT},
    {"HKEY_PERFORMANCE_TEXT", HKEY_PERFORMANCE_TEXT},
    {"HKEY_USERS", HKEY_USERS},
};

const std::map<DWORD, std::string> kRegistryTypes = {
    {REG_BINARY, "REG_BINARY"},
    {REG_DWORD, "REG_DWORD"},
    {REG_DWORD_BIG_ENDIAN, "REG_DWORD_BIG_ENDIAN"},
    {REG_EXPAND_SZ, "REG_EXPAND_SZ"},
    {REG_LINK, "REG_LINK"},
    {REG_MULTI_SZ, "REG_MULTI_SZ"},
    {REG_NONE, "REG_NONE"},
    {REG_QWORD, "REG_QWORD"},
    {REG_SZ, "REG_SZ"},
    {REG_FULL_RESOURCE_DESCRIPTOR, "REG_FULL_RESOURCE_DESCRIPTOR"},
    {REG_RESOURCE_LIST, "REG_RESOURCE_LIST"},
};

Status getUsernameFromKey(const std::string& key, std::string& rUsername) {
  if (!boost::starts_with(key, "HKEY_USERS")) {
    return Status(1, "Can not extract username from non-HKEY_USERS key");
  }

  auto toks = osquery::split(key, kRegSep);
  if (toks.size() < 2) {
    return Status(
        1, "Improperly-formatted HKEY_USERS key, cannot extract username");
  }

  PSID sid;
  if (!ConvertStringSidToSidA(toks[1].c_str(), &sid)) {
    return Status(GetLastError(), "Could not convert string to sid");
  } else {
    wchar_t accntName[UNLEN] = {0};
    wchar_t domName[DNLEN] = {0};
    unsigned long accntNameLen = UNLEN;
    unsigned long domNameLen = DNLEN;
    SID_NAME_USE eUse;
    if (!LookupAccountSidW(nullptr,
                           sid,
                           accntName,
                           &accntNameLen,
                           domName,
                           &domNameLen,
                           &eUse)) {
      return Status(GetLastError(), "Could not find sid");
    } else {
      rUsername = std::move(wstringToString(accntName));
    }
  }
  return Status(0, "OK");
}

inline void explodeRegistryPath(const std::string& path,
                                std::string& rHive,
                                std::string& rKey) {
  auto toks = osquery::split(path, kRegSep);
  rHive = toks.front();
  toks.erase(toks.begin());
  rKey = osquery::join(toks, kRegSep);
}

/// Microsoft helper function for getting the contents of a registry key
void queryKey(const std::string& keyPath, QueryData& results) {
  std::string hive;
  std::string key;
  explodeRegistryPath(keyPath, hive, key);

  if (kRegistryHives.count(hive) != 1) {
    return;
  }

  HKEY hRegistryHandle;
  auto ret = RegOpenKeyEx(kRegistryHives.at(hive),
                          TEXT(key.c_str()),
                          0,
                          KEY_READ,
                          &hRegistryHandle);

  if (ret != ERROR_SUCCESS) {
    return;
  }

  const DWORD maxKeyLength = 255;
  const DWORD maxValueName = 16383;
  TCHAR achClass[MAX_PATH] = TEXT("");
  DWORD cchClassName = MAX_PATH;
  DWORD cSubKeys = 0;
  DWORD cbMaxSubKey;
  DWORD cchMaxClass;
  DWORD cValues;
  DWORD cchMaxValueName;
  DWORD cbMaxValueData;
  DWORD cbSecurityDescriptor;
  DWORD retCode;
  FILETIME ftLastWriteTime;
  retCode = RegQueryInfoKey(hRegistryHandle,
                            achClass,
                            &cchClassName,
                            nullptr,
                            &cSubKeys,
                            &cbMaxSubKey,
                            &cchMaxClass,
                            &cValues,
                            &cchMaxValueName,
                            &cbMaxValueData,
                            &cbSecurityDescriptor,
                            &ftLastWriteTime);

  TCHAR achKey[maxKeyLength];
  DWORD cbName;

  // Process registry subkeys
  if (cSubKeys > 0) {
    for (DWORD i = 0; i < cSubKeys; i++) {
      cbName = maxKeyLength;
      retCode = RegEnumKeyEx(hRegistryHandle,
                             i,
                             achKey,
                             &cbName,
                             nullptr,
                             nullptr,
                             nullptr,
                             &ftLastWriteTime);
      if (retCode != ERROR_SUCCESS) {
        continue;
      }
      Row r;
      r["key"] = keyPath;
      r["type"] = "subkey";
      r["name"] = achKey;
      r["path"] = keyPath + kRegSep + achKey;
      r["mtime"] = std::to_string(osquery::filetimeToUnixtime(ftLastWriteTime));
      results.push_back(r);
    }
  }

  if (cValues <= 0) {
    return;
  }

  DWORD cchValue = maxKeyLength;
  TCHAR achValue[maxValueName];
  BYTE* bpDataBuff =
      (cbMaxValueData == 0) ? nullptr : new BYTE[cbMaxValueData]();

  // Process registry values
  for (size_t i = 0, retCode = ERROR_SUCCESS; i < cValues; i++) {
    size_t cnt = 0;
    cchValue = maxValueName;
    achValue[0] = '\0';

    retCode = RegEnumValue(hRegistryHandle,
                           static_cast<DWORD>(i),
                           achValue,
                           &cchValue,
                           nullptr,
                           nullptr,
                           nullptr,
                           nullptr);
    if (retCode != ERROR_SUCCESS) {
      continue;
    }

    DWORD lpData = cbMaxValueData;
    DWORD lpType;

    retCode = RegQueryValueEx(
        hRegistryHandle, achValue, 0, &lpType, bpDataBuff, &lpData);
    if (retCode != ERROR_SUCCESS) {
      continue;
    }

    // It's possible for registry entries to have been inserted incorrectly
    // resulting in non-null-terminated strings
    if (bpDataBuff != nullptr && lpData != 0 &&
        kRegistryStringTypes.find(lpType) != kRegistryStringTypes.end()) {
      bpDataBuff[lpData - 1] = 0x00;
    }

    Row r;
    r["key"] = keyPath;
    r["name"] = ((achValue[0] == '\0') ? "(Default)" : achValue);
    r["path"] = keyPath + kRegSep + achValue;
    if (kRegistryTypes.count(lpType) > 0) {
      r["type"] = kRegistryTypes.at(lpType);
    } else {
      r["type"] = "UNKNOWN";
    }
    r["mtime"] = std::to_string(osquery::filetimeToUnixtime(ftLastWriteTime));

    if (bpDataBuff != nullptr) {
      /// REG_LINK is a Unicode string, which in Windows is wchar_t
      char* regLinkStr = nullptr;
      if (lpType == REG_LINK) {
        regLinkStr = new char[cbMaxValueData];
        const size_t newSize = cbMaxValueData;
        size_t convertedChars = 0;
        wcstombs_s(&convertedChars,
                   regLinkStr,
                   newSize,
                   (wchar_t*)bpDataBuff,
                   _TRUNCATE);
      }

      BYTE* bpDataBuffTmp = bpDataBuff;
      std::vector<std::string> multiSzStrs;
      std::vector<char> regBinary;
      std::string data;

      switch (lpType) {
      case REG_FULL_RESOURCE_DESCRIPTOR:
      case REG_RESOURCE_LIST:
      case REG_BINARY:
        for (size_t i = 0; i < cbMaxValueData; i++) {
          regBinary.push_back((char)bpDataBuff[i]);
        }
        boost::algorithm::hex(
            regBinary.begin(), regBinary.end(), std::back_inserter(data));
        r["data"] = data;
        break;
      case REG_DWORD:
        r["data"] = std::to_string(*((int*)bpDataBuff));
        break;
      case REG_DWORD_BIG_ENDIAN:
        r["data"] = std::to_string(_byteswap_ulong(*((int*)bpDataBuff)));
        break;
      case REG_EXPAND_SZ:
        r["data"] = std::string((char*)bpDataBuff);
        break;
      case REG_LINK:
        r["data"] = std::string(regLinkStr);
        break;
      case REG_MULTI_SZ:
        while (*bpDataBuffTmp != 0x00) {
          std::string s((char*)bpDataBuffTmp);
          bpDataBuffTmp += s.size() + 1;
          multiSzStrs.push_back(s);
        }
        r["data"] = boost::algorithm::join(multiSzStrs, ",");
        break;
      case REG_NONE:
        r["data"] = "(zero-length binary value)";
        break;
      case REG_QWORD:
        r["data"] = std::to_string(*((unsigned long long*)bpDataBuff));
        break;
      case REG_SZ:
        r["data"] = std::string((char*)bpDataBuff);
        break;
      default:
        r["data"] = "";
        break;
      }
      if (regLinkStr != nullptr) {
        delete[](regLinkStr);
      }
      ZeroMemory(bpDataBuff, cbMaxValueData);
    }
    results.push_back(r);
  }
  if (bpDataBuff != nullptr) {
    delete[](bpDataBuff);
  }
  RegCloseKey(hRegistryHandle);
}

static inline void populateDefaultKeys(std::set<std::string>& rKeys) {
  boost::copy(kRegistryHives | boost::adaptors::map_keys,
              std::inserter(rKeys, rKeys.end()));
}

static inline void populateSubkeys(std::set<std::string>& rKeys,
                                   bool replaceKeys = false) {
  std::set<std::string> newKeys;
  if (!replaceKeys) {
    newKeys = rKeys;
  }

  for (const auto& key : rKeys) {
    QueryData regResults;
    queryKey(key, regResults);
    for (const auto& r : regResults) {
      if (r.at("type") == "subkey") {
        newKeys.insert(r.at("path"));
      }
    }
  }
  rKeys = std::move(newKeys);
}

static inline void appendSubkeyToKeys(const std::string& subkey,
                                      std::set<std::string>& rKeys) {
  std::set<std::string> newKeys{};
  for (auto& key : rKeys) {
    newKeys.insert(std::move(key) + kRegSep + subkey);
  }
  rKeys = std::move(newKeys);
}

static inline Status populateAllKeysRecursive(
    std::set<std::string>& rKeys,
    size_t currDepth = 1,
    size_t maxDepth = kRegMaxRecursiveDepth) {
  if (currDepth > maxDepth) {
    return Status(1, "Max recursive depth reached");
  }

  auto size_pre = rKeys.size();
  populateSubkeys(rKeys);
  if (size_pre < rKeys.size()) {
    auto status = populateAllKeysRecursive(rKeys, ++currDepth);
    if (!status.ok()) {
      return status;
    }
  }

  return Status(0, "OK");
}

Status expandRegistryGlobs(const std::string& pattern,
                           std::set<std::string>& results) {
  auto pathElems = osquery::split(pattern, kRegSep);
  if (pathElems.size() == 0) {
    return Status(0, "OK");
  }

  /*
   * Pattern is '%%', grab everything.
   * Note that if '%%' is present but not at the end of the pattern,
   * then it is treated like a single glob.
   */
  if (boost::ends_with(pathElems[0], kSQLGlobRecursive) &&
      pathElems.size() == 1) {
    populateDefaultKeys(results);
    return populateAllKeysRecursive(results);
  }

  // Special handling to insert default keys when glob present in first elem
  if (pathElems[0].find(kSQLGlobWildcard) != std::string::npos) {
    populateDefaultKeys(results);
    pathElems.erase(pathElems.begin());
  } else {
    results.insert(pathElems[0]);
    pathElems.erase(pathElems.begin());
  }

  for (auto&& elem = pathElems.begin(); elem != pathElems.end(); ++elem) {
    // We only care about  a recursive glob if it comes at the end of the
    // pattern i.e. 'HKEY_LOCAL_MACHINE\SOFTWARE\%%'
    if (boost::ends_with(*elem, kSQLGlobRecursive) &&
        *elem == pathElems.back()) {
      return populateAllKeysRecursive(results);
    } else if ((*elem).find(kSQLGlobWildcard) != std::string::npos) {
      populateSubkeys(results, true);
    } else {
      appendSubkeyToKeys(*elem, results);
    }
  }
  return Status(0, "OK");
}

static inline void maybeWarnLocalUsers(const std::set<std::string>& rKeys) {
  std::string hive, _;
  for (const auto& key : rKeys) {
    explodeRegistryPath(key, hive, _);
    if (hive == "HKEY_CURRENT_USER" ||
        hive == "HKEY_CURRENT_USER_LOCAL_SETTINGS") {
      LOG(WARNING) << "CURRENT_USER hives are not queryable by osqueryd; "
                      "query HKEY_USERS with the desired users SID instead";
      break;
    }
  }
}

QueryData genRegistry(QueryContext& context) {
  QueryData results;
  std::set<std::string> keys;

  if (!(context.hasConstraint("key", EQUALS) ||
        context.hasConstraint("key", LIKE))) {
    // We default to display all HIVEs
    expandRegistryGlobs(kSQLGlobWildcard, keys);
  } else {
    keys = context.constraints["key"].getAll(EQUALS);
    auto status = context.expandConstraints(
        "key",
        LIKE,
        keys,
        ([&](const std::string& pattern, std::set<std::string>& out) {
          std::set<std::string> resolvedKeys;
          auto status = expandRegistryGlobs(pattern, resolvedKeys);
          out.insert(resolvedKeys.begin(), resolvedKeys.end());
          return status;
        }));
    if (!status.ok()) {
      LOG(INFO) << "Failed to expand globs: " + status.getMessage();
    }
  }

  maybeWarnLocalUsers(keys);

  for (const auto& key : keys) {
    queryKey(key, results);
  }
  return results;
}
}
}
<|endoftext|>