Function wrapping implementation and Symtab extensions for undefined symbols
[dyninst.git] / valueAdded / wrapper / src / wrapper.C
1 /*
2  * wrapper.C
3  *
4  * Function wrapping mutator for Dyninst
5  *
6  * This program implements a simple function wrapping mutator for Dyninst.
7  * 
8  * Caveats: 
9  *   * Currently works only with the binary rewriter. 
10  *   * The wrapper must call the original function using a different
11  *     name, so that we can determine the mangled name to rename the
12  *     original to. If the wrapper will not call the original, use
13  *     function replacement instead.
14  *   * Requires that the specified names match single functions, so users
15  *     must use mangled names if functions are overloaded. 
16  *
17  * Usage: the program takes the following arguments:
18  *   -i <input executable>  : the file containing functions to be wrapped.
19  *   -o <output executable> : the file that will be produced. Should be
20  *                            used instead of the original input.
21  *   -l <library>           : file containing wrapper functions. May be repeated.
22  *   -s <specfile>          : specifies wrapping actions. Each line must contain
23  *                            three space-separated names:
24  *                              original wrapper clone
25  *                            where original is the function to be wrapped, 
26  *                            wrapper is the name of the wrapper function, 
27  *                            and clone is the new name we give original to allow
28  *                            it to be called directly. 
29  *                            These names may be mangled or demangled, but must
30  *                            follow the SymtabAPI naming specifications for each.
31  */
32
33 #include <iostream>
34 #include <string.h>
35 #include <unistd.h>
36 #include <sys/types.h>
37 #include <sys/stat.h>
38 #include <fstream>
39
40 #include "BPatch.h"
41 #include "BPatch_function.h"
42 #include "BPatch_Vector.h"
43 #include "BPatch_thread.h"
44 #include "BPatch_snippet.h"
45 #include "BPatch_process.h"
46 #include "BPatch_binaryEdit.h"
47 #include "BPatch_callbacks.h"
48 #include "BPatch_basicBlock.h"
49 #include "BPatch_flowGraph.h"
50 #include "CodeSource.h"
51 #include "PatchCFG.h"
52 #include "CFG.h"
53
54 #include <map>
55
56 using namespace std;
57 using namespace Dyninst;
58 using namespace SymtabAPI;
59
60 typedef vector<BPatch_function *> FuncVec;
61 typedef vector<Symbol *> SymbolVec;
62
63 BPatch bpatch;
64 BPatch_binaryEdit *app = NULL;
65 string output;
66
67 vector<string> specFiles;
68
69 void usage() {
70    cerr << "Usage: wrapper <arguments>" << endl;
71    cerr << "\t -i <input file name>" << endl;
72    cerr << "\t -o <output file name>" << endl;
73    cerr << "\t -l <wrapper library> (may be specified multiple times)" << endl;
74    cerr << "\t -s <specification file> (may be specified multiple times)" << endl;
75 }
76
77 // Look up the Symbol representing the clone. We use this instead of 
78 // a string directly so that we can use the compiler name mangler. 
79 Symbol *findCloneSymbol(Symtab *symtab, string clone) {
80    // Look up the provided name in each opened Symtab;
81    // we're looking for an undefined symbol.
82    Symbol *sym = NULL;
83    SymbolVec candidates;
84
85    if (symtab->findSymbol(candidates, clone, Symbol::ST_UNKNOWN,
86                           anyName, false, false, true)) {
87       // We can get multiple hits if the symbol is in both the static and dynamic 
88       // symbol tables. SymtabAPI doesn't have a filter mechanism (yet...) so
89       // go through it now.
90       for (SymbolVec::iterator c = candidates.begin(); c != candidates.end(); ++c) {
91          if ((*c)->isInDynSymtab()) {
92             if (sym && sym->isInDynSymtab()) {
93                cerr << "Error: multiple candidates found for clone " << clone << endl;
94                exit(0);
95             }
96          }
97          // Take any option we have, just prefer the dynamic symbol table one. 
98          sym = *c;
99       }
100       if (candidates.empty()) {
101          cerr << "Error: internal error in SymtabAPI" << endl;
102          exit(0);
103       }
104    }
105    if (!sym) {
106       cerr << "Error: could not find symbol for clone " << clone << endl;
107       exit(0);
108    }
109    return sym;
110 }
111
112 // Do the work of wrapping a function: look up the BPatch_function objects
113 // for the original and wrapper, the Symbol for the clone, and call wrapFunction.
114 //
115 // Requires that there be a single match for each. 
116 bool wrap(string original, string wrapper, string clone) {
117    FuncVec o;
118    app->getImage()->findFunction(original.c_str(), o);
119    if (o.empty()) return false;
120    if (o.size() > 1) { 
121       cerr << "Error: multiple matches found for original function " << original << endl;
122       exit(0);
123    }
124
125
126    FuncVec w;
127    app->getImage()->findFunction(wrapper.c_str(), w, false, false, true);
128    if (w.empty()) return false;
129    if (w.size() > 1) {
130       cerr << "Error: multiple matches found for wrapper function " << wrapper << endl;
131       exit(0);
132    }
133
134    ParseAPI::CodeObject *co = w[0]->getPatchAPIFunc()->function()->obj();
135    ParseAPI::SymtabCodeSource *src = dynamic_cast<ParseAPI::SymtabCodeSource *>(co->cs());
136    if (!src) {
137       cerr << "Error: wrapper function created from non-SymtabAPI code source" << endl;
138       exit(0);
139    }
140
141    // Find the symbol for the clone
142    Symbol *c = findCloneSymbol(src->getSymtabObject(), clone);
143    if (!c) exit(0);
144
145    return app->wrapFunction(o[0], w[0], c);
146 }
147
148 // Simple option parser
149
150 bool parseOptions(int argc, char **argv) {
151    bool outputSet = false;
152
153    int c;
154    while ((c = getopt(argc, argv, "i:o:l:s:h")) != -1) {
155       switch(c) {
156          case 'i':
157             if (app) {
158                cerr << "Error: must specify only one input file" << endl;
159                exit(0);
160             }
161             app = bpatch.openBinary(optarg, false);
162             break;
163          case 'o':
164             outputSet = true;
165             output = optarg;
166             break;
167          case 'l': {
168             if (!app) {
169                cerr << "Error: must specify input file before loading wrapper libraries" << endl;
170                exit(0);
171             }
172             app->loadLibrary(optarg);
173             break;
174          }
175          case 's':
176             specFiles.push_back(optarg);
177             break;
178          case 'h':
179             usage();
180             exit(0);
181             break;
182          default:
183             cerr << "Unknown argument " << (char) c << endl;
184             usage();
185             exit(0);
186             break;
187       }
188    }
189    if (!outputSet) {
190       cerr << "Error: no output file name given" << endl;
191       usage();
192       exit(0);
193    }
194    if (!app) {
195       cerr << "Error: no input file name given" << endl;
196       usage();
197       exit(0);
198    }
199    if (specFiles.empty()) {
200       cerr << "Warning: no spec files for wrapping provided" << endl;
201    }
202
203    return true;
204 }
205
206 // And read and parse the spec file. Each line must have the format
207 // original wrapper clone
208 // and we're pretty dumb about how we handle it. 
209 void handleSpecFile(string specfile) {
210    ifstream file;
211    file.open(specfile.c_str());
212    if (!file.is_open()) {
213       cerr << "Error: failed to open specfile " << specfile << endl;
214       exit(0);
215    }
216    while (file.good()) {
217       string original, wrapper, clone;
218       file >> original;
219       file >> wrapper;
220       file >> clone;
221       // We get a trailing blank line...
222       if (original == "" || wrapper == "" || clone == "") continue;      
223
224       wrap(original, wrapper, clone);
225    }
226    file.close();
227 }
228
229
230 int main(int argc, char** argv)
231 {    
232    parseOptions(argc, argv);
233
234    for (vector<string>::iterator iter = specFiles.begin(); 
235         iter != specFiles.end(); ++iter) {
236       handleSpecFile(*iter);
237    }
238    
239    app->writeFile(output.c_str());
240    
241    return 1;
242 }