msi: Fix buffer size calculation in get_keypath.
[wine] / dlls / msi / string.c
1 /*
2  * String Table Functions
3  *
4  * Copyright 2002-2004, Mike McCormack for CodeWeavers
5  * Copyright 2007 Robert Shearman for CodeWeavers
6  * Copyright 2010 Hans Leidekker for CodeWeavers
7  *
8  * This library is free software; you can redistribute it and/or
9  * modify it under the terms of the GNU Lesser General Public
10  * License as published by the Free Software Foundation; either
11  * version 2.1 of the License, or (at your option) any later version.
12  *
13  * This library is distributed in the hope that it will be useful,
14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
16  * Lesser General Public License for more details.
17  *
18  * You should have received a copy of the GNU Lesser General Public
19  * License along with this library; if not, write to the Free Software
20  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
21  */
22
23 #define COBJMACROS
24
25 #include <stdarg.h>
26 #include <assert.h>
27
28 #include "windef.h"
29 #include "winbase.h"
30 #include "winerror.h"
31 #include "wine/debug.h"
32 #include "wine/unicode.h"
33 #include "msi.h"
34 #include "msiquery.h"
35 #include "objbase.h"
36 #include "objidl.h"
37 #include "msipriv.h"
38 #include "winnls.h"
39
40 #include "query.h"
41
42 WINE_DEFAULT_DEBUG_CHANNEL(msidb);
43
44 typedef struct _msistring
45 {
46     USHORT persistent_refcount;
47     USHORT nonpersistent_refcount;
48     LPWSTR str;
49 } msistring;
50
51 struct string_table
52 {
53     UINT maxcount;         /* the number of strings */
54     UINT freeslot;
55     UINT codepage;
56     UINT sortcount;
57     msistring *strings; /* an array of strings */
58     UINT *sorted;       /* index */
59 };
60
61 static BOOL validate_codepage( UINT codepage )
62 {
63     if (codepage != CP_ACP && !IsValidCodePage( codepage ))
64     {
65         WARN("invalid codepage %u\n", codepage);
66         return FALSE;
67     }
68     return TRUE;
69 }
70
71 static string_table *init_stringtable( int entries, UINT codepage )
72 {
73     string_table *st;
74
75     if (!validate_codepage( codepage ))
76         return NULL;
77
78     st = msi_alloc( sizeof (string_table) );
79     if( !st )
80         return NULL;    
81     if( entries < 1 )
82         entries = 1;
83
84     st->strings = msi_alloc_zero( sizeof (msistring) * entries );
85     if( !st->strings )
86     {
87         msi_free( st );
88         return NULL;    
89     }
90
91     st->sorted = msi_alloc( sizeof (UINT) * entries );
92     if( !st->sorted )
93     {
94         msi_free( st->strings );
95         msi_free( st );
96         return NULL;
97     }
98
99     st->maxcount = entries;
100     st->freeslot = 1;
101     st->codepage = codepage;
102     st->sortcount = 0;
103
104     return st;
105 }
106
107 VOID msi_destroy_stringtable( string_table *st )
108 {
109     UINT i;
110
111     for( i=0; i<st->maxcount; i++ )
112     {
113         if( st->strings[i].persistent_refcount ||
114             st->strings[i].nonpersistent_refcount )
115             msi_free( st->strings[i].str );
116     }
117     msi_free( st->strings );
118     msi_free( st->sorted );
119     msi_free( st );
120 }
121
122 static int st_find_free_entry( string_table *st )
123 {
124     UINT i, sz, *s;
125     msistring *p;
126
127     TRACE("%p\n", st);
128
129     if( st->freeslot )
130     {
131         for( i = st->freeslot; i < st->maxcount; i++ )
132             if( !st->strings[i].persistent_refcount &&
133                 !st->strings[i].nonpersistent_refcount )
134                 return i;
135     }
136     for( i = 1; i < st->maxcount; i++ )
137         if( !st->strings[i].persistent_refcount &&
138             !st->strings[i].nonpersistent_refcount )
139             return i;
140
141     /* dynamically resize */
142     sz = st->maxcount + 1 + st->maxcount/2;
143     p = msi_realloc_zero( st->strings, sz*sizeof(msistring) );
144     if( !p )
145         return -1;
146
147     s = msi_realloc( st->sorted, sz*sizeof(UINT) );
148     if( !s )
149     {
150         msi_free( p );
151         return -1;
152     }
153
154     st->strings = p;
155     st->sorted = s;
156
157     st->freeslot = st->maxcount;
158     st->maxcount = sz;
159     if( st->strings[st->freeslot].persistent_refcount ||
160         st->strings[st->freeslot].nonpersistent_refcount )
161         ERR("oops. expected freeslot to be free...\n");
162     return st->freeslot;
163 }
164
165 static int find_insert_index( const string_table *st, UINT string_id )
166 {
167     int i, c, low = 0, high = st->sortcount - 1;
168
169     while (low <= high)
170     {
171         i = (low + high) / 2;
172         c = strcmpW( st->strings[string_id].str, st->strings[st->sorted[i]].str );
173
174         if (c < 0)
175             high = i - 1;
176         else if (c > 0)
177             low = i + 1;
178         else
179             return -1; /* already exists */
180     }
181     return high + 1;
182 }
183
184 static void insert_string_sorted( string_table *st, UINT string_id )
185 {
186     int i;
187
188     i = find_insert_index( st, string_id );
189     if (i == -1)
190         return;
191
192     memmove( &st->sorted[i] + 1, &st->sorted[i], (st->sortcount - i) * sizeof(UINT) );
193     st->sorted[i] = string_id;
194     st->sortcount++;
195 }
196
197 static void set_st_entry( string_table *st, UINT n, LPWSTR str, USHORT refcount, enum StringPersistence persistence )
198 {
199     if (persistence == StringPersistent)
200     {
201         st->strings[n].persistent_refcount = refcount;
202         st->strings[n].nonpersistent_refcount = 0;
203     }
204     else
205     {
206         st->strings[n].persistent_refcount = 0;
207         st->strings[n].nonpersistent_refcount = refcount;
208     }
209
210     st->strings[n].str = str;
211
212     insert_string_sorted( st, n );
213
214     if( n < st->maxcount )
215         st->freeslot = n + 1;
216 }
217
218 static UINT msi_string2idA( const string_table *st, LPCSTR buffer, UINT *id )
219 {
220     DWORD sz;
221     UINT r = ERROR_INVALID_PARAMETER;
222     LPWSTR str;
223
224     TRACE("Finding string %s in string table\n", debugstr_a(buffer) );
225
226     if( buffer[0] == 0 )
227     {
228         *id = 0;
229         return ERROR_SUCCESS;
230     }
231
232     sz = MultiByteToWideChar( st->codepage, 0, buffer, -1, NULL, 0 );
233     if( sz <= 0 )
234         return r;
235     str = msi_alloc( sz*sizeof(WCHAR) );
236     if( !str )
237         return ERROR_NOT_ENOUGH_MEMORY;
238     MultiByteToWideChar( st->codepage, 0, buffer, -1, str, sz );
239
240     r = msi_string2idW( st, str, id );
241     msi_free( str );
242
243     return r;
244 }
245
246 static int msi_addstring( string_table *st, UINT n, const CHAR *data, int len, USHORT refcount, enum StringPersistence persistence )
247 {
248     LPWSTR str;
249     int sz;
250
251     if( !data )
252         return 0;
253     if( !data[0] )
254         return 0;
255     if( n > 0 )
256     {
257         if( st->strings[n].persistent_refcount ||
258             st->strings[n].nonpersistent_refcount )
259             return -1;
260     }
261     else
262     {
263         if( ERROR_SUCCESS == msi_string2idA( st, data, &n ) )
264         {
265             if (persistence == StringPersistent)
266                 st->strings[n].persistent_refcount += refcount;
267             else
268                 st->strings[n].nonpersistent_refcount += refcount;
269             return n;
270         }
271         n = st_find_free_entry( st );
272         if( n == -1 )
273             return -1;
274     }
275
276     if( n < 1 )
277     {
278         ERR("invalid index adding %s (%d)\n", debugstr_a( data ), n );
279         return -1;
280     }
281
282     /* allocate a new string */
283     if( len < 0 )
284         len = strlen(data);
285     sz = MultiByteToWideChar( st->codepage, 0, data, len, NULL, 0 );
286     str = msi_alloc( (sz+1)*sizeof(WCHAR) );
287     if( !str )
288         return -1;
289     MultiByteToWideChar( st->codepage, 0, data, len, str, sz );
290     str[sz] = 0;
291
292     set_st_entry( st, n, str, refcount, persistence );
293
294     return n;
295 }
296
297 int msi_addstringW( string_table *st, const WCHAR *data, int len, USHORT refcount, enum StringPersistence persistence )
298 {
299     UINT n;
300     LPWSTR str;
301
302     if( !data )
303         return 0;
304     if( !data[0] )
305         return 0;
306
307     if( msi_string2idW( st, data, &n ) == ERROR_SUCCESS )
308     {
309         if (persistence == StringPersistent)
310             st->strings[n].persistent_refcount += refcount;
311         else
312             st->strings[n].nonpersistent_refcount += refcount;
313         return n;
314     }
315
316     n = st_find_free_entry( st );
317     if( n == -1 )
318         return -1;
319
320     /* allocate a new string */
321     if(len<0)
322         len = strlenW(data);
323     TRACE("%s, n = %d len = %d\n", debugstr_w(data), n, len );
324
325     str = msi_alloc( (len+1)*sizeof(WCHAR) );
326     if( !str )
327         return -1;
328     memcpy( str, data, len*sizeof(WCHAR) );
329     str[len] = 0;
330
331     set_st_entry( st, n, str, refcount, persistence );
332
333     return n;
334 }
335
336 /* find the string identified by an id - return null if there's none */
337 const WCHAR *msi_string_lookup_id( const string_table *st, UINT id )
338 {
339     if( id == 0 )
340         return szEmpty;
341
342     if( id >= st->maxcount )
343         return NULL;
344
345     if( id && !st->strings[id].persistent_refcount && !st->strings[id].nonpersistent_refcount)
346         return NULL;
347
348     return st->strings[id].str;
349 }
350
351 /*
352  *  msi_id2stringA
353  *
354  *  [in] st         - pointer to the string table
355  *  [in] id         - id of the string to retrieve
356  *  [out] buffer    - destination of the UTF8 string
357  *  [in/out] sz     - number of bytes available in the buffer on input
358  *                    number of bytes used on output
359  *
360  *   The size includes the terminating nul character.  Short buffers
361  *  will be filled, but not nul terminated.
362  */
363 static UINT msi_id2stringA( const string_table *st, UINT id, LPSTR buffer, UINT *sz )
364 {
365     UINT len;
366     const WCHAR *str;
367     int n;
368
369     TRACE("Finding string %d of %d\n", id, st->maxcount);
370
371     str = msi_string_lookup_id( st, id );
372     if( !str )
373         return ERROR_FUNCTION_FAILED;
374
375     len = WideCharToMultiByte( st->codepage, 0, str, -1, NULL, 0, NULL, NULL );
376
377     if( !buffer )
378     {
379         *sz = len;
380         return ERROR_SUCCESS;
381     }
382
383     if( len > *sz )
384     {
385         n = strlenW( str ) + 1;
386         while( n && (len > *sz) )
387             len = WideCharToMultiByte( st->codepage, 0, 
388                            str, --n, NULL, 0, NULL, NULL );
389     }
390     else
391         n = -1;
392
393     *sz = WideCharToMultiByte( st->codepage, 0, str, n, buffer, len, NULL, NULL );
394
395     return ERROR_SUCCESS;
396 }
397
398 /*
399  *  msi_string2idW
400  *
401  *  [in] st         - pointer to the string table
402  *  [in] str        - string to find in the string table
403  *  [out] id        - id of the string, if found
404  */
405 UINT msi_string2idW( const string_table *st, LPCWSTR str, UINT *id )
406 {
407     int i, c, low = 0, high = st->sortcount - 1;
408
409     while (low <= high)
410     {
411         i = (low + high) / 2;
412         c = strcmpW( str, st->strings[st->sorted[i]].str );
413
414         if (c < 0)
415             high = i - 1;
416         else if (c > 0)
417             low = i + 1;
418         else
419         {
420             *id = st->sorted[i];
421             return ERROR_SUCCESS;
422         }
423     }
424
425     return ERROR_INVALID_PARAMETER;
426 }
427
428 static void string_totalsize( const string_table *st, UINT *datasize, UINT *poolsize )
429 {
430     UINT i, len, holesize;
431
432     if( st->strings[0].str || st->strings[0].persistent_refcount || st->strings[0].nonpersistent_refcount)
433         ERR("oops. element 0 has a string\n");
434
435     *poolsize = 4;
436     *datasize = 0;
437     holesize = 0;
438     for( i=1; i<st->maxcount; i++ )
439     {
440         if( !st->strings[i].persistent_refcount )
441         {
442             TRACE("[%u] nonpersistent = %s\n", i, debugstr_w(st->strings[i].str));
443             (*poolsize) += 4;
444         }
445         else if( st->strings[i].str )
446         {
447             TRACE("[%u] = %s\n", i, debugstr_w(st->strings[i].str));
448             len = WideCharToMultiByte( st->codepage, 0,
449                      st->strings[i].str, -1, NULL, 0, NULL, NULL);
450             if( len )
451                 len--;
452             (*datasize) += len;
453             if (len>0xffff)
454                 (*poolsize) += 4;
455             (*poolsize) += holesize + 4;
456             holesize = 0;
457         }
458         else
459             holesize += 4;
460     }
461     TRACE("data %u pool %u codepage %x\n", *datasize, *poolsize, st->codepage );
462 }
463
464 static const WCHAR szStringData[] = {
465     '_','S','t','r','i','n','g','D','a','t','a',0 };
466 static const WCHAR szStringPool[] = {
467     '_','S','t','r','i','n','g','P','o','o','l',0 };
468
469 HRESULT msi_init_string_table( IStorage *stg )
470 {
471     USHORT zero[2] = { 0, 0 };
472     UINT ret;
473
474     /* create the StringPool stream... add the zero string to it*/
475     ret = write_stream_data(stg, szStringPool, zero, sizeof zero, TRUE);
476     if (ret != ERROR_SUCCESS)
477         return E_FAIL;
478
479     /* create the StringData stream... make it zero length */
480     ret = write_stream_data(stg, szStringData, NULL, 0, TRUE);
481     if (ret != ERROR_SUCCESS)
482         return E_FAIL;
483
484     return S_OK;
485 }
486
487 string_table *msi_load_string_table( IStorage *stg, UINT *bytes_per_strref )
488 {
489     string_table *st = NULL;
490     CHAR *data = NULL;
491     USHORT *pool = NULL;
492     UINT r, datasize = 0, poolsize = 0, codepage;
493     DWORD i, count, offset, len, n, refs;
494
495     r = read_stream_data( stg, szStringPool, TRUE, (BYTE **)&pool, &poolsize );
496     if( r != ERROR_SUCCESS)
497         goto end;
498     r = read_stream_data( stg, szStringData, TRUE, (BYTE **)&data, &datasize );
499     if( r != ERROR_SUCCESS)
500         goto end;
501
502     if ( (poolsize > 4) && (pool[1] & 0x8000) )
503         *bytes_per_strref = LONG_STR_BYTES;
504     else
505         *bytes_per_strref = sizeof(USHORT);
506
507     count = poolsize/4;
508     if( poolsize > 4 )
509         codepage = pool[0] | ( (pool[1] & ~0x8000) << 16 );
510     else
511         codepage = CP_ACP;
512     st = init_stringtable( count, codepage );
513     if (!st)
514         goto end;
515
516     offset = 0;
517     n = 1;
518     i = 1;
519     while( i<count )
520     {
521         /* the string reference count is always the second word */
522         refs = pool[i*2+1];
523
524         /* empty entries have two zeros, still have a string id */
525         if (pool[i*2] == 0 && refs == 0)
526         {
527             i++;
528             n++;
529             continue;
530         }
531
532         /*
533          * If a string is over 64k, the previous string entry is made null
534          * and its the high word of the length is inserted in the null string's
535          * reference count field.
536          */
537         if( pool[i*2] == 0)
538         {
539             len = (pool[i*2+3] << 16) + pool[i*2+2];
540             i += 2;
541         }
542         else
543         {
544             len = pool[i*2];
545             i += 1;
546         }
547
548         if ( (offset + len) > datasize )
549         {
550             ERR("string table corrupt?\n");
551             break;
552         }
553
554         r = msi_addstring( st, n, data+offset, len, refs, StringPersistent );
555         if( r != n )
556             ERR("Failed to add string %d\n", n );
557         n++;
558         offset += len;
559     }
560
561     if ( datasize != offset )
562         ERR("string table load failed! (%08x != %08x), please report\n", datasize, offset );
563
564     TRACE("Loaded %d strings\n", count);
565
566 end:
567     msi_free( pool );
568     msi_free( data );
569
570     return st;
571 }
572
573 UINT msi_save_string_table( const string_table *st, IStorage *storage, UINT *bytes_per_strref )
574 {
575     UINT i, datasize = 0, poolsize = 0, sz, used, r, codepage, n;
576     UINT ret = ERROR_FUNCTION_FAILED;
577     CHAR *data = NULL;
578     USHORT *pool = NULL;
579
580     TRACE("\n");
581
582     /* construct the new table in memory first */
583     string_totalsize( st, &datasize, &poolsize );
584
585     TRACE("%u %u %u\n", st->maxcount, datasize, poolsize );
586
587     pool = msi_alloc( poolsize );
588     if( ! pool )
589     {
590         WARN("Failed to alloc pool %d bytes\n", poolsize );
591         goto err;
592     }
593     data = msi_alloc( datasize );
594     if( ! data )
595     {
596         WARN("Failed to alloc data %d bytes\n", poolsize );
597         goto err;
598     }
599
600     used = 0;
601     codepage = st->codepage;
602     pool[0] = codepage & 0xffff;
603     pool[1] = codepage >> 16;
604     if (st->maxcount > 0xffff)
605     {
606         pool[1] |= 0x8000;
607         *bytes_per_strref = LONG_STR_BYTES;
608     }
609     else
610         *bytes_per_strref = sizeof(USHORT);
611
612     n = 1;
613     for( i=1; i<st->maxcount; i++ )
614     {
615         if( !st->strings[i].persistent_refcount )
616         {
617             pool[ n*2 ] = 0;
618             pool[ n*2 + 1] = 0;
619             n++;
620             continue;
621         }
622
623         sz = datasize - used;
624         r = msi_id2stringA( st, i, data+used, &sz );
625         if( r != ERROR_SUCCESS )
626         {
627             ERR("failed to fetch string\n");
628             sz = 0;
629         }
630         if( sz && (sz < (datasize - used ) ) )
631             sz--;
632
633         if (sz)
634             pool[ n*2 + 1 ] = st->strings[i].persistent_refcount;
635         else
636             pool[ n*2 + 1 ] = 0;
637         if (sz < 0x10000)
638         {
639             pool[ n*2 ] = sz;
640             n++;
641         }
642         else
643         {
644             pool[ n*2 ] = 0;
645             pool[ n*2 + 2 ] = sz&0xffff;
646             pool[ n*2 + 3 ] = (sz>>16);
647             n += 2;
648         }
649         used += sz;
650         if( used > datasize  )
651         {
652             ERR("oops overran %d >= %d\n", used, datasize);
653             goto err;
654         }
655     }
656
657     if( used != datasize )
658     {
659         ERR("oops used %d != datasize %d\n", used, datasize);
660         goto err;
661     }
662
663     /* write the streams */
664     r = write_stream_data( storage, szStringData, data, datasize, TRUE );
665     TRACE("Wrote StringData r=%08x\n", r);
666     if( r )
667         goto err;
668     r = write_stream_data( storage, szStringPool, pool, poolsize, TRUE );
669     TRACE("Wrote StringPool r=%08x\n", r);
670     if( r )
671         goto err;
672
673     ret = ERROR_SUCCESS;
674
675 err:
676     msi_free( data );
677     msi_free( pool );
678
679     return ret;
680 }
681
682 UINT msi_get_string_table_codepage( const string_table *st )
683 {
684     return st->codepage;
685 }
686
687 UINT msi_set_string_table_codepage( string_table *st, UINT codepage )
688 {
689     if (validate_codepage( codepage ))
690     {
691         st->codepage = codepage;
692         return ERROR_SUCCESS;
693     }
694     return ERROR_FUNCTION_FAILED;
695 }